feat(core): commitment tree management functions

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

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

zcash-haskell

@@ -1 +1 @@
-Subproject commit 62cda9cc15621dead6fbfd7a4944840408d69da4
+Subproject commit 20851a4e48f768a492796fb828f16ae9745931dc