From 8986847679037f59ff5cf49c5a5c24328690ef6b Mon Sep 17 00:00:00 2001 From: Rene Vergara Date: Fri, 8 Nov 2024 14:47:43 -0600 Subject: [PATCH] feat(core): commitment tree management functions --- src/Zenith/Tree.hs | 110 ++++++++++++++++++++++++++++++++++++++++++++- test/Spec.hs | 49 ++++++++++++++++++++ zcash-haskell | 2 +- 3 files changed, 159 insertions(+), 2 deletions(-) diff --git a/src/Zenith/Tree.hs b/src/Zenith/Tree.hs index e4b5802..b77d9b6 100644 --- a/src/Zenith/Tree.hs +++ b/src/Zenith/Tree.hs @@ -15,7 +15,13 @@ import Data.Maybe (fromJust, isNothing) import qualified GHC.Generics as GHC import qualified Generics.SOP as SOP import ZcashHaskell.Orchard (combineOrchardNodes, getOrchardNodeValue) -import ZcashHaskell.Types (MerklePath(..), OrchardFrontier(..), OrchardTree(..)) +import ZcashHaskell.Sapling (combineSaplingNodes, getSaplingNodeValue) +import ZcashHaskell.Types + ( MerklePath(..) + , OrchardFrontier(..) + , OrchardTree(..) + , SaplingTree(..) + ) type Level = Int8 @@ -33,6 +39,7 @@ class Node v where getLevel :: v -> Level getHash :: v -> HexString getPosition :: v -> Position + getIndex :: v -> Int64 isFull :: v -> Bool isMarked :: v -> Bool mkNode :: Level -> Position -> HexString -> v @@ -45,6 +52,14 @@ instance Measured OrchardCommitment OrchardNode where Nothing -> OrchardNode 0 (hexString "00") 0 True 0 False Just val -> OrchardNode p val 0 True i False +type SaplingCommitment = HexString + +instance Measured SaplingCommitment SaplingNode where + measure sc p i = + case getSaplingNodeValue (hexBytes sc) of + Nothing -> SaplingNode 0 (hexString "00") 0 True 0 False + Just val -> SaplingNode p val 0 True i False + data Tree v = EmptyLeaf | Leaf !v @@ -151,6 +166,98 @@ path pos (Branch s x y) = | otherwise = [] path _ _ = Nothing +getNotePosition :: Monoid v => Node v => Tree v -> Int64 -> Maybe Position +getNotePosition (Leaf x) i + | getIndex x == i = Just $ getPosition x + | otherwise = Nothing +getNotePosition (Branch _ x y) i + | getIndex (value x) >= i = getNotePosition x i + | getIndex (value y) >= i = getNotePosition y i + | otherwise = Nothing +getNotePosition _ _ = Nothing + +truncateTree :: Monoid v => Node v => Tree v -> Int64 -> Tree v +truncateTree (Branch s x y) i + | getLevel s == 1 && getIndex (value x) == i = branch x EmptyLeaf + | getLevel s == 1 && getIndex (value y) == i = branch x y + | getIndex (value x) >= i = + branch (truncateTree x i) (getEmptyRoot (getLevel s)) + | getIndex (value y) >= i = branch x (truncateTree y i) +truncateTree x _ = x + +data SaplingNode = SaplingNode + { sn_position :: !Position + , sn_value :: !HexString + , sn_level :: !Level + , sn_full :: !Bool + , sn_index :: !Int64 + , sn_mark :: !Bool + } deriving stock (Eq, GHC.Generic) + deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo) + deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct SaplingNode + +instance Semigroup SaplingNode where + (<>) x y = + case combineSaplingNodes (sn_level x) (sn_value x) (sn_value y) of + Nothing -> x + Just newHash -> + SaplingNode + (max (sn_position x) (sn_position y)) + newHash + (1 + sn_level x) + (sn_full x && sn_full y) + (max (sn_index x) (sn_index y)) + (sn_mark x || sn_mark y) + +instance Monoid SaplingNode where + mempty = SaplingNode 0 (hexString "00") 0 False 0 False + mappend = (<>) + +instance Node SaplingNode where + getLevel = sn_level + getHash = sn_value + getPosition = sn_position + getIndex = sn_index + isFull = sn_full + isMarked = sn_mark + mkNode l p v = SaplingNode p v l True 0 False + +instance Show SaplingNode where + show = show . sn_value + +saplingSize :: SaplingTree -> Int64 +saplingSize tree = + (if isNothing (st_left tree) + then 0 + else 1) + + (if isNothing (st_right tree) + then 0 + else 1) + + foldl + (\x (i, p) -> + case p of + Nothing -> x + 0 + Just _ -> x + 2 ^ i) + 0 + (zip [1 ..] $ st_parents tree) + +mkSaplingTree :: SaplingTree -> Tree SaplingNode +mkSaplingTree tree = + foldl + (\t (i, n) -> + case n of + Just n' -> prunedBranch i 0 n' <> t + Nothing -> t <> getEmptyRoot i) + leafRoot + (zip [1 ..] $ st_parents tree) + where + leafRoot = + case st_right tree of + Just r' -> leaf (fromJust $ st_left tree) (pos - 1) 0 <> leaf r' pos 0 + Nothing -> leaf (fromJust $ st_left tree) pos 0 <> EmptyLeaf + pos = fromIntegral $ saplingSize tree - 1 + +-- | Orchard data OrchardNode = OrchardNode { on_position :: !Position , on_value :: !HexString @@ -186,6 +293,7 @@ instance Node OrchardNode where getLevel = on_level getHash = on_value getPosition = on_position + getIndex = on_index isFull = on_full isMarked = on_mark mkNode l p v = OrchardNode p v l True 0 False diff --git a/test/Spec.hs b/test/Spec.hs index a668fae..b5335ea 100644 --- a/test/Spec.hs +++ b/test/Spec.hs @@ -425,6 +425,20 @@ main = do Just p1 -> getSaplingPathAnchor cmu1 p1 `shouldBe` getHash (value updatedTree) + it "Find position by index" $ do + let tree = + SaplingCommitmentTree $ + hexString + "01b4a6049100dd1aa53fd850c388e0ed3782b039448aa43de5fd41d9e2419b375b0114abe3b473ac2dcd9ec3c62ff5cffb683f2e584d70ce4c1c118d67a8c34f1d691f00000000014278a19b3777ee0ceb48c4d469ab0a95217de56f27a4eaf9f19a90244bcff05301692acd3dd617e7cb44a17026211a31899cf27ef5960d59b8323ba105c754965800016f641fd51f68efae0cb742182483d97054bca8fcfb8036b1c165acdb0b27fc1b0000015965da8d105e2d412c7e13f920e607461c8cbb67a38b3a667898d1eaa064ba5701be82b8379a88cce7629153e0b3d1bf826782d3e0c367d615d362b479ead41d6401e1dbd43a62ac601d82de803fc3b6485182854056c6577bae231882b19184b03e000000018f624406b4ebda9ee13c062063e16f65ab5d725b80645037be7760e4ca7887730000000000000000000000000000" + let cmu1 = + hexString + "238a75ea513eb330cee527d3fa5000cd52620f116919a33afb9ac78ee1f91c17" :: SaplingCommitment + case getSaplingTreeParts tree of + Nothing -> assertFailure "Failed to get tree parts" + Just t1 -> do + let newTree = mkSaplingTree t1 + let updatedTree = append newTree (cmu1, 4) + getNotePosition updatedTree 4 `shouldBe` Just 145762 describe "Orchard" $ do let cmx1 = hexString @@ -590,6 +604,41 @@ main = do Just p1 -> do getOrchardPathAnchor cmx2 p1 `shouldBe` getHash (value updatedTree) + it "Find position by index" $ do + let tree = + OrchardCommitmentTree $ + hexString + "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000" + let cmx1 = + hexString + "1712ead46028d4349e234abf59e94e0640fe7a0829e2e2e17e1a931631810400" :: OrchardCommitment + let cmx2 = + hexString + "39f5ad39817fb432fa07c5feb3a957189fbe7662a4b5555ca95093b6d853cf07" :: OrchardCommitment + case getOrchardTreeParts tree of + Nothing -> assertFailure "Failed to get tree parts" + Just t1 -> do + let newTree = mkOrchardTree t1 + let updatedTree = foldl append newTree [(cmx1, 4), (cmx2, 5)] + getNotePosition updatedTree 4 `shouldBe` Just 39734 + it "Truncate tree" $ do + let tree = + OrchardCommitmentTree $ + hexString + "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000" + let cmx1 = + hexString + "1712ead46028d4349e234abf59e94e0640fe7a0829e2e2e17e1a931631810400" :: OrchardCommitment + let cmx2 = + hexString + "39f5ad39817fb432fa07c5feb3a957189fbe7662a4b5555ca95093b6d853cf07" :: OrchardCommitment + case getOrchardTreeParts tree of + Nothing -> assertFailure "Failed to get tree parts" + Just t1 -> do + let newTree = mkOrchardTree t1 + let updatedTree = foldl append newTree [(cmx1, 4), (cmx2, 5)] + let truncTree = truncateTree updatedTree 4 + getIndex (value truncTree) `shouldBe` 4 describe "Creating Tx" $ do describe "Full" $ do it "To Orchard" $ do diff --git a/zcash-haskell b/zcash-haskell index 62cda9c..20851a4 160000 --- a/zcash-haskell +++ b/zcash-haskell @@ -1 +1 @@ -Subproject commit 62cda9cc15621dead6fbfd7a4944840408d69da4 +Subproject commit 20851a4e48f768a492796fb828f16ae9745931dc