Test walkMerkle, walkMerkleV2, streamMerkle

2024-09-30 06:05:29 +04:00 · 2024-09-30 06:05:29 +04:00 · e29b15f90c
parent 1715e1dd92
commit e29b15f90c
5 changed files with 297 additions and 0 deletions
--- a/hbs2-core/hbs2-core.cabal
+++ b/hbs2-core/hbs2-core.cabal
@ -95,6 +95,7 @@ library
                    , HBS2.Hash
                    , HBS2.Merkle
                    , HBS2.Merkle.MetaData
+                    , HBS2.Merkle.Walk
                    , HBS2.Net.Auth.Schema
                    , HBS2.Net.Auth.GroupKeyAsymm
                    , HBS2.Net.Auth.GroupKeySymm
@ -186,6 +187,7 @@ library
                    , stm-chans
                    , string-conversions
                    , streaming
+                    , streaming-bytestring
                    , string-conversions
                    , suckless-conf
                    , template-haskell
--- a/hbs2-core/lib/HBS2/Merkle/Walk.hs
+++ b/hbs2-core/lib/HBS2/Merkle/Walk.hs
@ -0,0 +1,190 @@
+module HBS2.Merkle.Walk where
+
+import Codec.Serialise (DeserialiseFailure, deserialiseOrFail, serialise)
+import Control.Exception
+import Control.Monad
+import Control.Monad.Except
+import Control.Monad.Fix
+import Control.Monad.Trans.Class
+import Data.ByteString.Lazy qualified as BSL
+import Data.Foldable
+import Data.Functor
+import Data.String.Conversions (cs)
+import GHC.Generics (Generic)
+import Prettyprinter
+import Streaming (Of (..), Stream)
+import Streaming qualified as S
+import Streaming.ByteString qualified as Q
+import Streaming.Prelude qualified as S
+
+import HBS2.Base58
+import HBS2.Data.Types.Refs
+import HBS2.Hash
+import HBS2.Merkle
+
+walkMerkleV2
+    :: forall a m
+     . (Serialise (MTree a), Serialise (MTreeAnn a), Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> (Either WalkMerkleError a -> m ())
+    -> Hash HbSync
+    -> m ()
+walkMerkleV2 flookup sink =
+    walkMerkleV2' flookup \case
+        (Right (MLeaf s)) -> sink (Right s)
+        (Right (MNode _ _)) -> pure ()
+        Left e -> sink (Left e)
+  where
+    walkMerkleV2'
+        :: forall a m
+         . (Serialise (MTree a), Serialise (MTreeAnn a), Monad m)
+        => (Hash HbSync -> m (Maybe BSL.ByteString))
+        -> (Either WalkMerkleError (MTree a) -> m ())
+        -> Hash HbSync
+        -> m ()
+
+    walkMerkleV2' flookup sink root =
+        either (sink . Left) pure =<< runExceptT (go root)
+      where
+        go :: Hash HbSync -> ExceptT WalkMerkleError m ()
+        go = fix \go' hash -> do
+            bs <-
+                maybe (throwError $ MerkleHashNotFound hash) pure
+                    =<< lift (flookup hash)
+            either
+                (throwError . MerkleDeserialiseFailure hash)
+                (runWithTree (lift . sink . Right) (traverse_ go'))
+                ( (deserialiseOrFail @(MTree a) bs)
+                    <> (deserialiseOrFail bs <&> \(MTreeAnn {_mtaTree = t}) -> t)
+                )
+          where
+            runWithTree
+                :: forall m a
+                 . (Monad m)
+                => (MTree a -> m ())
+                -> ([Hash HbSync] -> m ())
+                -> MTree a
+                -> m ()
+            runWithTree h run = \case
+                n@(MLeaf _) -> h n
+                n@(MNode _ hashes) -> h n >> run hashes
+
+type WalkMerkleError = WalkMerkleError' (Hash HbSync)
+data WalkMerkleError' h
+    = MerkleHashNotFound h
+    | MerkleDeserialiseFailure h DeserialiseFailure
+    deriving (Generic)
+
+deriving instance Show (WalkMerkleError' (AsBase58 (Hash HbSync)))
+
+deriving via
+    (WalkMerkleError' (AsBase58 (Hash HbSync)))
+    instance
+        Show (WalkMerkleError' (Hash HbSync))
+
+instance Show (AsBase58 (Hash HbSync)) where
+    show (AsBase58 h) = show $ show $ pretty $ h
+
+instance Exception WalkMerkleError
+
+---
+
+streamMerkle
+    :: forall a m
+     . (Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of a) m (Either WalkMerkleError ())
+streamMerkle getB rv = (runExceptT . S.distribute) do
+    streamMerkle' getB rv
+
+streamMerkle'
+    :: forall a m
+     . (Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of a) (ExceptT WalkMerkleError m) ()
+streamMerkle' getB = do
+    walkMerkleV2 getB' \case
+        Left hashNotFound -> throwError hashNotFound
+        Right as -> S.each as
+  where
+    getB' :: Hash HbSync -> Stream (Of a) (ExceptT WalkMerkleError m) (Maybe BSL.ByteString)
+    getB' = lift . lift . getB
+
+streamCatFromMerkle
+    :: (Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of BSL.ByteString) m (Either WalkMerkleError ())
+streamCatFromMerkle getB =
+    fmap join
+        . (runExceptT . S.distribute)
+        . S.mapM
+            ( \(HashRef h) ->
+                maybe (throwError (MerkleHashNotFound h)) pure
+                    =<< (lift . getB) h
+            )
+        . streamMerkle (lift . getB)
+
+catFromMerkle
+    :: (Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> m (Either WalkMerkleError BSL.ByteString)
+catFromMerkle getB =
+    fmap (\(bs S.:> ehu) -> const bs <$> ehu)
+        . Q.toLazy
+        . Q.fromChunks
+        . (S.map cs . streamCatFromMerkle getB)
+
+---
+
+streamMerkle1
+    :: forall a m
+     . (Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of a) m (Either WalkMerkleError ())
+streamMerkle1 getB rv = (runExceptT . S.distribute) do
+    streamMerkle1' getB rv
+
+streamMerkle1'
+    :: forall a m
+     . (Serialise a, Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of a) (ExceptT WalkMerkleError m) ()
+streamMerkle1' getB hash = do
+    walkMerkle hash getB' \case
+        Left hashNotFound -> throwError (MerkleHashNotFound hashNotFound)
+        Right as -> S.each as
+  where
+    getB' :: Hash HbSync -> Stream (Of a) (ExceptT WalkMerkleError m) (Maybe BSL.ByteString)
+    getB' = lift . lift . getB
+
+streamCatFromMerkle1
+    :: (Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> Stream (Of BSL.ByteString) m (Either WalkMerkleError ())
+streamCatFromMerkle1 getB =
+    fmap join
+        . (runExceptT . S.distribute)
+        . S.mapM
+            ( \(HashRef h) ->
+                maybe (throwError (MerkleHashNotFound h)) pure
+                    =<< (lift . getB) h
+            )
+        . streamMerkle1 (lift . getB)
+
+catFromMerkle1
+    :: (Monad m)
+    => (Hash HbSync -> m (Maybe BSL.ByteString))
+    -> Hash HbSync
+    -> m (Either WalkMerkleError BSL.ByteString)
+catFromMerkle1 getB =
+    fmap (\(bs S.:> ehu) -> const bs <$> ehu)
+        . Q.toLazy
+        . Q.fromChunks
+        . (S.map cs . streamCatFromMerkle1 getB)
--- a/hbs2-peer/lib/HBS2/Peer/RPC/Client/RefChan.hs
+++ b/hbs2-peer/lib/HBS2/Peer/RPC/Client/RefChan.hs
@ -29,6 +29,7 @@ import UnliftIO
 getRefChanHeadHash :: forall proto m . ( MonadUnliftIO m
                                       , HasClientAPI RefChanAPI proto m
                                       , HasProtocol proto (ServiceProto RefChanAPI proto)
+
                                       )
                   => PubKey 'Sign 'HBS2Basic
                   -> m (Maybe HashRef)
--- a/hbs2-tests/hbs2-tests.cabal
+++ b/hbs2-tests/hbs2-tests.cabal
@ -1123,3 +1123,24 @@ executable test-refchan-notify
        , timeit


+executable test-walk-merkletree-cornercase
+    import:           shared-properties
+    import:           common-deps
+    default-language: Haskell2010
+    ghc-options:
+    -- other-extensions:
+    hs-source-dirs:   test
+    main-is:          TestWalkMerkleTreeCorner.hs
+    build-depends:
+          base, hbs2-core
+        , bytestring
+        , containers
+        , string-conversions
+        , interpolatedstring-perl6
+        -- , mtl
+        -- , mwc-random
+        -- , random
+        , safe
+        , serialise
+        , streaming
+        , text
--- a/hbs2-tests/test/TestWalkMerkleTreeCorner.hs
+++ b/hbs2-tests/test/TestWalkMerkleTreeCorner.hs
@ -0,0 +1,83 @@
+module Main where
+
+import Control.Monad
+import Control.Monad.State as State
+import Data.ByteString.Lazy (ByteString)
+import Data.List qualified as List
+import Data.Map (Map)
+import Data.Map qualified as Map
+import Data.String.Conversions (cs)
+
+import HBS2.Base58
+import HBS2.Data.Types.Refs
+import HBS2.Hash
+import HBS2.Merkle
+import HBS2.Merkle.Walk
+
+main :: IO ()
+main = do
+    -- Just "3.14"
+    print $ runMapStoreM do
+        hr <- mapStorePutBlock "3.14"
+        mapStoreReadBlock hr
+
+    print $ runMapStoreM do
+        h <- mapStorePutBlock "1853635"
+        (,)
+            -- Left (MerkleDeserialiseFailure "1337jkagZ5Knihh82JaXqzB3qAX4K298DuNp7jx5tGmw" (DeserialiseFailure 0 "expected list len"))
+            <$> catFromMerkle mapStoreReadBlock h
+            -- Right ""
+            <*> catFromMerkle1 mapStoreReadBlock h
+
+    -- (Right "123456789abcdefgh",Right "123456789abcdefgh")
+    -- (Left (MerkleHashNotFound "h3VBGX1u6JHDjR38z97xo6S3ruiynkc1kygGZgRVcit")
+    -- ,Left (MerkleHashNotFound "h3VBGX1u6JHDjR38z97xo6S3ruiynkc1kygGZgRVcit"))
+    mapM_ print $ runMapStoreM do
+        MerkleHash h <- storeAsMerkle mapStorePutBlock do
+            ((cs . (: [])) <$> "123456789abcdefgh") :: [ByteString]
+        r1 <- (,) <$> catFromMerkle mapStoreReadBlock h <*> catFromMerkle1 mapStoreReadBlock h
+
+        mapStoreDeleteBlock . (!! 2) =<< State.gets Map.keys
+        r2 <- (,) <$> catFromMerkle mapStoreReadBlock h <*> catFromMerkle1 mapStoreReadBlock h
+
+        pure [r1, r2]
+
+type MapStore = Map (Hash HbSync) ByteString
+newtype MapStoreM a = MapStoreM {unMapStoreM :: State MapStore a}
+    deriving newtype (Functor, Applicative, Monad, MonadState MapStore)
+
+runMapStoreM :: MapStoreM a -> a
+runMapStoreM = flip evalState mempty . unMapStoreM
+
+mapStorePutBlock :: ByteString -> MapStoreM (Hash HbSync)
+mapStorePutBlock bs =
+    h <$ State.modify (Map.insert h bs)
+  where
+    h = hashObject bs
+
+mapStoreReadBlock :: (Hash HbSync) -> MapStoreM (Maybe ByteString)
+mapStoreReadBlock h =
+    State.gets (Map.lookup h)
+
+mapStoreDeleteBlock :: (Hash HbSync) -> MapStoreM ()
+mapStoreDeleteBlock h =
+    State.modify (Map.delete h)
+
+---
+
+class (Monad m) => StoreAsMerkle m h b a where
+    storeAsMerkle :: (b -> m h) -> a -> m MerkleHash
+
+instance StoreAsMerkle MapStoreM (Hash HbSync) ByteString [ByteString] where
+    storeAsMerkle = \putB bs -> do
+        hashes <- mapM putB bs
+        storeAsMerkle putB (HashRef <$> hashes)
+
+instance StoreAsMerkle MapStoreM (Hash HbSync) ByteString [HashRef] where
+    storeAsMerkle = \putB hrs -> do
+        let
+            treeChildNum = 3
+            hashListChunk = 2
+            ptree = toPTree (MaxSize hashListChunk) (MaxNum treeChildNum) hrs
+        MerkleHash <$> do
+            makeMerkle 0 ptree \(_, _, bs) -> (void . putB) bs