From d2c52508d1932dabf9fb9cd406a12b644b17c8b5 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Tue, 19 Nov 2024 15:03:36 -0600
Subject: [PATCH] feat: batch append commitments

---
 CHANGELOG.md       |  2 ++
 src/Zenith/Core.hs | 16 ++++++++++++----
 src/Zenith/Tree.hs | 27 +++++++++++++++++++++++++++
 test/Spec.hs       |  5 +++--
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index abbf77e..25ee291 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Function `prepareTxV2` implementing `PrivacyPolicy`
 - Functionality to shield transparent balance
 - Functionality to de-shield shielded notes
+- Native commitment trees
+	- Batch append to trees in O(log n)
 
 ### Changed
 
diff --git a/src/Zenith/Core.hs b/src/Zenith/Core.hs
index 43bfe6b..6b0cdab 100644
--- a/src/Zenith/Core.hs
+++ b/src/Zenith/Core.hs
@@ -1369,7 +1369,9 @@ updateCommitmentTrees pool zHost zPort zNet = do
             return InvalidTree
           Just t1 -> do
             let newTree = mkSaplingTree t1
-            return $ foldl' append newTree saplingComm
+            let zippedSapComms =
+                  zip [(getPosition (value newTree) + 1) ..] saplingComm
+            return $ batchAppend newTree zippedSapComms
       Just (sTree, sSync) -> do
         logDebugN $ ">Sapling tree found, synced to " <> T.pack (show sSync)
         saplingNotes <- liftIO $ getShieldedOutputs pool (sSync + 1) zNet
@@ -1380,7 +1382,9 @@ updateCommitmentTrees pool zHost zPort zNet = do
                    , fromSqlKey (entityKey y)))
                 saplingNotes
         logDebugN ">got shielded outputs"
-        return $ foldl' append sTree saplingComm
+        let zippedSapComms =
+              zip [(getPosition (value sTree) + 1) ..] saplingComm
+        return $ batchAppend sTree zippedSapComms
   newOrchTree <-
     case oTdb of
       Nothing -> do
@@ -1402,7 +1406,9 @@ updateCommitmentTrees pool zHost zPort zNet = do
             return InvalidTree
           Just t1 -> do
             let newTree = mkOrchardTree t1
-            return $ foldl' append newTree orchardComm
+            let zippedOrchComms =
+                  zip [(getPosition (value newTree) + 1) ..] orchardComm
+            return $ batchAppend newTree zippedOrchComms
       Just (oTree, oSync) -> do
         logDebugN $ ">Orchard tree found, synced to " <> T.pack (show oSync)
         orchardNotes <- liftIO $ getOrchardActions pool (oSync + 1) zNet
@@ -1413,7 +1419,9 @@ updateCommitmentTrees pool zHost zPort zNet = do
                    , fromSqlKey (entityKey y)))
                 orchardNotes
         logDebugN ">got orchard actions"
-        return $ foldl' append oTree orchardComm
+        let zippedOrchComms =
+              zip [(getPosition (value oTree) + 1) ..] orchardComm
+        return $ batchAppend oTree zippedOrchComms
   case newSapTree of
     Branch {} -> do
       logInfoN ">Saving updated Sapling tree to db"
diff --git a/src/Zenith/Tree.hs b/src/Zenith/Tree.hs
index 9e70530..1a52243 100644
--- a/src/Zenith/Tree.hs
+++ b/src/Zenith/Tree.hs
@@ -225,6 +225,33 @@ countLeaves (Leaf _) = 1
 countLeaves EmptyLeaf = 0
 countLeaves InvalidTree = 0
 
+batchAppend ::
+     Measured a v
+  => Node v => Monoid v => Tree v -> [(Int32, (a, Int64))] -> Tree v
+batchAppend x [] = x
+batchAppend (Branch s x y) notes
+  | isFull s = InvalidTree
+  | isFull (value x) = branch x (batchAppend y notes)
+  | otherwise =
+    branch
+      (batchAppend x (take leftSide notes))
+      (batchAppend y (drop leftSide notes))
+  where
+    leftSide = fromIntegral $ 2 ^ getLevel (value x) - countLeaves x
+batchAppend (PrunedBranch k) notes
+  | isFull k = InvalidTree
+  | otherwise =
+    branch
+      (batchAppend (getEmptyRoot (getLevel k - 1)) (take leftSide notes))
+      (batchAppend (getEmptyRoot (getLevel k - 1)) (drop leftSide notes))
+  where
+    leftSide = fromIntegral $ 2 ^ (getLevel k - 1)
+batchAppend EmptyLeaf notes
+  | length notes == 1 =
+    leaf (fst $ snd $ head notes) (fst $ head notes) (snd $ snd $ head notes)
+  | otherwise = InvalidTree
+batchAppend _ notes = InvalidTree
+
 data SaplingNode = SaplingNode
   { sn_position :: !Position
   , sn_value :: !HexString
diff --git a/test/Spec.hs b/test/Spec.hs
index 6b044fc..ca66599 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -663,7 +663,7 @@ main = do
         it "Validate large load" $ do
           pool <- runNoLoggingT $ initPool "/home/rav/Zenith/zenith.db"
           maxBlock <- getMaxBlock pool $ ZcashNetDB TestNet
-          let startBlock = maxBlock - 2000
+          let startBlock = maxBlock - 310000
           zebraTreesIn <-
             getCommitmentTrees
               pool
@@ -690,7 +690,8 @@ main = do
                          ( getHex $ orchActionCmx $ entityVal y
                          , fromSqlKey $ entityKey y))
                       oAct
-              let updatedTree = foldl' append newTree cmxs
+              let posCmx = zip [(getPosition (value newTree) + 1) ..] cmxs
+              let updatedTree = batchAppend newTree posCmx
               let finalAnchor =
                     getOrchardTreeAnchor $
                     OrchardCommitmentTree $ ztiOrchard zebraTreesOut