Map query.

IntersectMBO · May 4, 2021 · 01fa978 · 01fa978
1 parent 3c62579
commit 01fa978
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Query.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Ledger/Query.hs
@@ -39,7 +39,9 @@ import qualified Data.Map.Strict as Map
 
 -- | Different queries supported by the ledger for all block types, indexed
 -- by the result type.
-data Query blk result = BlockQuery (BlockQuery blk result)
+newtype Query blk result = BlockQuery
+  { unQuery :: BlockQuery blk result
+  }
 
 instance (ShowProxy (BlockQuery blk)) => ShowProxy (Query blk) where
   showProxy (Proxy :: Proxy (Query blk)) = "Query (" ++ showProxy (Proxy @(BlockQuery blk)) ++ ")"

diff --git a/ouroboros-network/src/Ouroboros/Network/Protocol/LocalStateQuery/Client.hs b/ouroboros-network/src/Ouroboros/Network/Protocol/LocalStateQuery/Client.hs
@@ -22,6 +22,14 @@ module Ouroboros.Network.Protocol.LocalStateQuery.Client (
 
       -- * Utilities
     , mapLocalStateQueryClient
+
+    , mapQueryOnClientStAcquired
+    , mapQueryOnClientStAcquiring
+    , mapQueryOnClientStIdle
+    , mapQueryOnClientStQuerying
+
+    , mapQueryOnLocalStateQueryClient
+
     , Some (..)
     ) where
 
@@ -34,12 +42,16 @@ import           Network.TypedProtocol.Core
 import           Ouroboros.Network.Protocol.LocalStateQuery.Type
 import           Ouroboros.Network.Protocol.LocalStateQuery.Codec (Some(..))
 
-
 newtype LocalStateQueryClient block point (query :: Type -> Type) m a =
     LocalStateQueryClient {
       runLocalStateQueryClient :: m (ClientStIdle block point query m a)
     }
 
+mapQueryOnLocalStateQueryClient :: Monad m => (forall r. query r -> query' r) -> LocalStateQueryClient block point query m a -> LocalStateQueryClient block point query' m a
+mapQueryOnLocalStateQueryClient f localStateQueryClient = LocalStateQueryClient
+  { runLocalStateQueryClient = mapQueryOnClientStIdle f <$> runLocalStateQueryClient localStateQueryClient
+  }
+
 localStateQueryClientNull :: MonadTimer m => LocalStateQueryClient block point query m a
 localStateQueryClientNull =
     LocalStateQueryClient $ forever $ threadDelay 43200 {- day in seconds -}
@@ -59,6 +71,11 @@ data ClientStIdle block point query (m :: Type -> Type) a where
   SendMsgDone    :: a
                  -> ClientStIdle      block point query m a
 
+mapQueryOnClientStIdle :: Monad m => (forall r . query r -> query' r) -> ClientStIdle block point query m a -> ClientStIdle block point query' m a
+mapQueryOnClientStIdle f client = case client of
+  SendMsgAcquire mPoint clientStAcquiring -> SendMsgAcquire mPoint (mapQueryOnClientStAcquiring f clientStAcquiring)
+  SendMsgDone a -> SendMsgDone a
+
 -- | In the 'StAcquiring' protocol state, the client does not have agency.
 -- Instead it is waiting for:
 --
@@ -74,6 +91,12 @@ data ClientStAcquiring block point query m a = ClientStAcquiring {
                       -> m (ClientStIdle  block point query m a)
     }
 
+mapQueryOnClientStAcquiring :: Monad m => (forall r . query r -> query' r) -> ClientStAcquiring block point query m a -> ClientStAcquiring block point query' m a
+mapQueryOnClientStAcquiring f acquiring = ClientStAcquiring
+  { recvMsgAcquired = fmap (mapQueryOnClientStAcquired f) (recvMsgAcquired acquiring)
+  , recvMsgFailure = fmap (mapQueryOnClientStIdle f) . recvMsgFailure acquiring
+  }
+
 -- | In the 'StAcquired' protocol state, the client has agency and must send:
 --
 --  * a query
@@ -92,6 +115,13 @@ data ClientStAcquired block point query m a where
   SendMsgRelease   :: m (ClientStIdle   block point query m a)
                    -> ClientStAcquired  block point query m a
 
+mapQueryOnClientStAcquired :: Monad m => (forall r . query r -> query' r) -> ClientStAcquired block point query m a -> ClientStAcquired block point query' m a
+mapQueryOnClientStAcquired f clientStAcquired = case clientStAcquired of
+  SendMsgQuery queryResult clientStQuerying -> SendMsgQuery (f queryResult) (mapQueryOnClientStQuerying f clientStQuerying)
+  SendMsgReAcquire mPoint clientStAcquiring -> SendMsgReAcquire mPoint (mapQueryOnClientStAcquiring f clientStAcquiring)
+  SendMsgRelease runClientStIdle -> SendMsgRelease (mapQueryOnClientStIdle f <$> runClientStIdle)
+
+
 -- | In the 'StQuerying' protocol state, the client does not have agency.
 -- Instead it is waiting for:
 --
@@ -101,6 +131,11 @@ data ClientStQuerying block point query m a result = ClientStQuerying {
       recvMsgResult :: result -> m (ClientStAcquired block point query m a)
     }
 
+mapQueryOnClientStQuerying :: Monad m => (forall r. query r -> query' r) -> ClientStQuerying block point query m a result -> ClientStQuerying block point query' m a result
+mapQueryOnClientStQuerying f clientStQuerying = ClientStQuerying
+  { recvMsgResult = fmap (mapQueryOnClientStAcquired f) . recvMsgResult clientStQuerying
+  }
+
 -- | Transform a 'LocalStateQueryClient' by mapping over the query and query
 -- result values.
 --