walkMerkleConditional, streamMerkleConditional

hbs2-core/lib/HBS2/Merkle/Walk.hs

@@ -1,11 +1,12 @@
 module HBS2.Merkle.Walk where
 
-import Codec.Serialise (DeserialiseFailure, deserialiseOrFail, serialise)
+import Codec.Serialise
 import Control.Exception
 import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Fix
 import Control.Monad.Trans.Class
+import Data.Bool
 import Data.ByteString.Lazy qualified as BSL
 import Data.Foldable
 import Data.Functor
@@ -39,44 +40,65 @@ walkMerkleV2
     -> Hash HbSync
     -> m ()
 walkMerkleV2 flookup sink =
-    walkMerkleV2' flookup \case
+    either (sink . Left) pure <=< runExceptT . go
-        (Right (MLeaf s)) -> sink (Right s)
-        (Right (MNode _ _)) -> pure ()
-        Left e -> sink (Left e)
   where
-    walkMerkleV2'
+    go :: Hash HbSync -> ExceptT WalkMerkleError m ()
-        :: forall a m
+    go = fix \go' hash -> do
-         . (Serialise (MTree a), Serialise (MTreeAnn a), Monad m)
+        bs <-
-        => (Hash HbSync -> m (Maybe BSL.ByteString))
+            maybe (throwError $ MerkleHashNotFound hash) pure
-        -> (Either WalkMerkleError (MTree a) -> m ())
+                =<< lift (flookup hash)
-        -> Hash HbSync
+        either
-        -> m ()
+            (throwError . MerkleDeserialiseFailure hash)
+            (runWithTree (withMLeaf (lift . sink . Right)) (traverse_ go'))
+            ( (deserialiseOrFail @(MTree a) bs)
+                <> (deserialiseOrFail bs <&> \(MTreeAnn {_mtaTree = t}) -> t)
+            )
 
-    walkMerkleV2' flookup sink root =
+data WalkStep a
-        either (sink . Left) pure =<< runExceptT (go root)
+    = WalkLeaf a
-      where
+    | WalkProcessedTree (Hash HbSync)
-        go :: Hash HbSync -> ExceptT WalkMerkleError m ()
+    | WalkSkippedTree (Hash HbSync)
-        go = fix \go' hash -> do
+
-            bs <-
+walkMerkleConditional
-                maybe (throwError $ MerkleHashNotFound hash) pure
+    :: forall a m
-                    =<< lift (flookup hash)
+     . (Serialise (MTree a), Serialise (MTreeAnn a), Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> (Hash HbSync -> m Bool)
+    -> (Either WalkMerkleError (WalkStep a) -> m ())
+    -> Hash HbSync
+    -> m ()
+walkMerkleConditional getB p sink =
+    either (sink . Left) pure <=< runExceptT . go
+  where
+    go :: Hash HbSync -> ExceptT WalkMerkleError m ()
+    go = fix \go' h ->
+        (lift . p) h >>= bool (sinkRight (WalkSkippedTree h)) do
+            bs <- maybe (throwError $ MerkleHashNotFound h) pure =<< (lift . getB) h
             either
-                (throwError . MerkleDeserialiseFailure hash)
+                (throwError . MerkleDeserialiseFailure h)
-                (runWithTree (lift . sink . Right) (traverse_ go'))
+                (runWithTree (withMLeaf (sinkRight . WalkLeaf)) (traverse_ go'))
                 ( (deserialiseOrFail @(MTree a) bs)
                     <> (deserialiseOrFail bs <&> \(MTreeAnn {_mtaTree = t}) -> t)
                 )
-          where
+            sinkRight (WalkProcessedTree h)
-            runWithTree
+      where
-                :: forall m a
+        sinkRight = lift . sink . Right
-                 . (Monad m)
+
-                => (MTree a -> m ())
+withMLeaf :: (Applicative m) => (a -> m ()) -> MTree a -> m ()
-                -> ([Hash HbSync] -> m ())
+withMLeaf f = \case
-                -> MTree a
+    MLeaf s -> f s
-                -> m ()
+    MNode _ _ -> pure ()
-            runWithTree h run = \case
+
-                n@(MLeaf _) -> h n
+runWithTree
-                n@(MNode _ hashes) -> h n >> run hashes
+    :: forall m a
+     . (Monad m)
+    => (MTree a -> m ())
+    -> ([Hash HbSync] -> m ())
+    -> MTree a
+    -> m ()
+runWithTree f run = \case
+    n@(MLeaf _) -> f n
+    n@(MNode _ hashes) -> f n >> run hashes
 
 type WalkMerkleError = WalkMerkleError' (Hash HbSync)
 data WalkMerkleError' h
@@ -98,6 +120,32 @@ instance Exception WalkMerkleError
 
 ---
 
+streamMerkleConditionalEither
+    :: forall a m
+     . (Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> (Hash HbSync -> m Bool)
+    -> Hash HbSync
+    -> Stream (Of (WalkStep a)) m (Either WalkMerkleError ())
+streamMerkleConditionalEither getB p h =
+    (runExceptT . S.distribute) do
+        streamMerkleConditional (lift . getB) (lift . p) h
+
+streamMerkleConditional
+    :: forall a m
+     . ( Serialise a
+       , Monad m
+       , MonadError WalkMerkleError m
+       )
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> (Hash HbSync -> m Bool)
+    -> Hash HbSync
+    -> Stream (Of (WalkStep a)) m ()
+streamMerkleConditional getB p = do
+    walkMerkleConditional (lift . getB) (lift . p) \case
+        Left hashNotFound -> throwError hashNotFound
+        Right as -> S.yield as
+
 streamMerkle
     :: forall a m
      . (Serialise a, Monad m)
@@ -107,6 +155,7 @@ streamMerkle
 streamMerkle getB rv = (runExceptT . S.distribute) do
     streamMerkle' getB rv
 
+-- | Работает для деревьев, у которых в узлах лежат [a]
 streamMerkle'
     :: forall a m
      . (Serialise a, Monad m)
@@ -114,7 +163,7 @@ streamMerkle'
     -> Hash HbSync
     -> Stream (Of a) (ExceptT WalkMerkleError m) ()
 streamMerkle' getB = do
-    walkMerkleV2 getB' \case
+    walkMerkleV2 @[a] getB' \case
         Left hashNotFound -> throwError hashNotFound
         Right as -> S.each as
   where