From 6073392af1d196dc5f0f07cdc2352385c0af8394 Mon Sep 17 00:00:00 2001
From: Georgy Lukyanov <mail@geo2a.info>
Date: Mon, 10 Feb 2025 13:14:37 +0000
Subject: [PATCH] Refine local-state-query.cddl by specifying possible queries

---
 .../cddl/specs/local-state-query.cddl         | 11 +++++++-
 .../Protocol/LocalStateQuery/Codec/CDDL.hs    | 25 ++++++++++++++++---
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/ouroboros-network-protocols/cddl/specs/local-state-query.cddl b/ouroboros-network-protocols/cddl/specs/local-state-query.cddl
index 991ee9c1ac..15d4e8c2b8 100644
--- a/ouroboros-network-protocols/cddl/specs/local-state-query.cddl
+++ b/ouroboros-network-protocols/cddl/specs/local-state-query.cddl
@@ -18,7 +18,16 @@ acquireFailurePointNotOnChain = 1
 failure      = acquireFailurePointTooOld
              / acquireFailurePointNotOnChain
 
-query        = any
+blockQuery      = [0, point]
+getSystemStart  = [1]
+getChainBlockNo = [2]
+getChainPoint   = [3]
+
+query        = blockQuery
+             / getSystemStart
+             / getChainBlockNo
+             / getChainPoint
+
 result       = any
 
 msgAcquire   = [0, point]
diff --git a/ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/LocalStateQuery/Codec/CDDL.hs b/ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/LocalStateQuery/Codec/CDDL.hs
index b4a0a899c7..1373658102 100644
--- a/ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/LocalStateQuery/Codec/CDDL.hs
+++ b/ouroboros-network-protocols/testlib/Ouroboros/Network/Protocol/LocalStateQuery/Codec/CDDL.hs
@@ -1,4 +1,6 @@
 {-# LANGUAGE DeriveGeneric              #-}
+{-# LANGUAGE DeriveAnyClass             #-}
+{-# LANGUAGE DerivingStrategies         #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GADTs                      #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -20,14 +22,29 @@ import Ouroboros.Network.Protocol.BlockFetch.Codec.CDDL (Block, BlockPoint)
 import Ouroboros.Network.Protocol.LocalStateQuery.Codec
 import Ouroboros.Network.Protocol.LocalStateQuery.Type
 import Test.Data.CDDL (Any)
-import Test.QuickCheck (Arbitrary (..))
+import Test.QuickCheck (Arbitrary (..), oneof)
 
 newtype Result = Result Any
-  deriving (Eq, Show, Arbitrary, Serialise, Generic, NFData)
+  deriving stock (Eq, Show, Generic)
+  deriving newtype (Arbitrary, Serialise, NFData)
+
+data QueryPayload =
+    BlockQuery Any
+  | GetSystemStart
+  | GetChainBlockNo
+  | GetChainPoint
+  deriving stock (Eq, Show, Generic)
+  deriving anyclass (Serialise, NFData)
+
+instance Arbitrary QueryPayload where
+    arbitrary = oneof [ BlockQuery <$> arbitrary
+                      , pure GetSystemStart
+                      , pure GetChainBlockNo
+                      , pure GetChainPoint
+                      ]
 
--- TODO: add payload to the query
 data Query result where
-    Query :: Any -> Query Result
+    Query :: QueryPayload -> Query Result
 
 instance NFData (Query result) where
   rnf (Query a) = rnf a