Restore Merkle tree pruning for the wallet (#353)

* pruning fix

* it works

* more things

* everything ready

* everything ready

* final touches

* changelog

* clippy

* final error found

* no removing paths

* modified extend

* remove optimization

* comment addressed

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 ### Added
+- [\#353](https://github.com/Manta-Network/manta-rs/pull/353) Restore Merkle tree pruning for the wallet.
 
 ### Changed
 

manta-crypto/src/merkle_tree/forest.rs

@@ -787,15 +787,14 @@ where
     }
 }
 
-impl<C, T, M, L, const N: usize> Rollback
-    for MerkleForest<C, TreeArray<C, ForkedTree<C, T, M, L>, N>>
+impl<C, M, L, const N: usize> Rollback
+    for MerkleForest<C, TreeArray<C, ForkedTree<C, Partial<C, M, L>, M, L>, N>>
 where
     C: Configuration + ?Sized,
     C::Index: FixedIndex<N>,
-    T: Tree<C>,
     M: Default + InnerMap<C>,
-    L: LeafMap<C> + Default,
-    LeafDigest<C>: Clone + Default,
+    L: Default + LeafMap<C>,
+    LeafDigest<C>: Clone + Default + PartialEq,
     InnerDigest<C>: Clone + Default + PartialEq,
 {
     #[inline]
@@ -808,7 +807,7 @@ where
     #[inline]
     fn commit(&mut self) {
         for tree in self.forest.as_mut() {
-            tree.merge_fork(&self.parameters);
+            tree.merge_fork_partial(&self.parameters);
         }
     }
 }

manta-crypto/src/merkle_tree/fork.rs

@@ -19,7 +19,7 @@
 use crate::merkle_tree::{
     capacity,
     inner_tree::{BTreeMap, InnerMap, PartialInnerTree},
-    leaf_map::{LeafMap, TrivialLeafVec},
+    leaf_map::{LeafBTreeMap, LeafMap},
     partial::Partial,
     path::{CurrentInnerPath, InnerPath},
     Configuration, CurrentPath, InnerDigest, Leaf, LeafDigest, Node, Parameters, Parity, Path,
@@ -274,7 +274,7 @@ impl Default for BaseContribution {
     Hash(bound = "InnerDigest<C>: Hash, M: Hash, L: Hash"),
     PartialEq(bound = "InnerDigest<C>: PartialEq, M: PartialEq, L: PartialEq")
 )]
-struct Branch<C, M = BTreeMap<C>, L = TrivialLeafVec<C>>
+struct Branch<C, M = BTreeMap<C>, L = LeafBTreeMap<C>>
 where
     C: Configuration + ?Sized,
     M: Default + InnerMap<C>,
@@ -397,6 +397,17 @@ where
         leaf_digests
     }
 
+    /// Extracts the non-base leaves with their markings from `base_contribution` and `data`.
+    #[inline]
+    fn extract_leaves_with_markings(
+        base_contribution: BaseContribution,
+        data: Partial<C, M, L>,
+    ) -> Vec<(bool, LeafDigest<C>)> {
+        let mut marked_digests = data.into_leaves_with_markings();
+        mem::drop(marked_digests.drain(0..base_contribution as usize));
+        marked_digests
+    }
+
     /// Tries to rebase `self` at `base`.
     #[inline]
     fn try_rebase<T>(&mut self, parameters: &Parameters<C>, base: &T) -> bool
@@ -579,6 +590,31 @@ where
             "Should have been able to extend extracted leaves."
         );
     }
+
+    /// Merges `self` into the [`Partial`] tree `base`, keeping the markings.
+    ///
+    /// # Panics
+    ///
+    /// This method panics if the [`Tree::extend_digests`] method returns an `Err` variant because
+    /// the capacity invariant should have prevented the addition of leaves to this branch if they
+    /// would have exceeded the capacity limit of `base`.
+    #[inline]
+    fn merge_partial(self, parameters: &Parameters<C>, base: &mut Partial<C, M, L>)
+    where
+        L: Default,
+        M: Default,
+        InnerDigest<C>: Clone + Default + PartialEq,
+        LeafDigest<C>: Clone + Default,
+    {
+        assert!(
+            base.extend_with_marked_digests(
+                parameters,
+                Self::extract_leaves_with_markings(self.base_contribution, self.data)
+            )
+            .is_ok(),
+            "Should have been able to extend extracted leaves."
+        );
+    }
 }
 
 /// Merkle Tree Fork
@@ -587,7 +623,7 @@ where
     Debug(bound = "P::Weak: Debug, L: Debug, InnerDigest<C>: Debug, M: Debug"),
     Default(bound = "L: Default, InnerDigest<C>: Default")
 )]
-pub struct Fork<C, T, P = pointer::SingleThreaded, M = BTreeMap<C>, L = TrivialLeafVec<C>>
+pub struct Fork<C, T, P = pointer::SingleThreaded, M = BTreeMap<C>, L = LeafBTreeMap<C>>
 where
     C: Configuration + ?Sized,
     T: Tree<C>,
@@ -839,7 +875,7 @@ where
     Hash(bound = "T: Hash, L: Hash, InnerDigest<C>: Hash, M: Hash"),
     PartialEq(bound = "T: PartialEq, L: PartialEq, InnerDigest<C>: PartialEq, M: PartialEq")
 )]
-pub struct ForkedTree<C, T, M = BTreeMap<C>, L = TrivialLeafVec<C>>
+pub struct ForkedTree<C, T, M = BTreeMap<C>, L = LeafBTreeMap<C>>
 where
     C: Configuration + ?Sized,
     T: Tree<C>,
@@ -994,9 +1030,11 @@ where
     }
 }
 
-impl<C> ForkedTree<C, Partial<C>>
+impl<C, M, L> ForkedTree<C, Partial<C, M, L>, M, L>
 where
     C: Configuration + ?Sized,
+    M: Default + InnerMap<C>,
+    L: Default + LeafMap<C>,
     LeafDigest<C>: Clone + Default + PartialEq,
     InnerDigest<C>: Clone + Default + PartialEq,
 {
@@ -1013,6 +1051,13 @@ where
             parameters,
         )
     }
+
+    /// Merges the fork of the base partial tree back into the trunk.
+    #[inline]
+    pub fn merge_fork_partial(&mut self, parameters: &Parameters<C>) {
+        mem::take(&mut self.branch).merge_partial(parameters, &mut self.base);
+        self.reset_fork(parameters)
+    }
 }
 
 impl<C, T, M, L> Tree<C> for ForkedTree<C, T, M, L>
@@ -1072,7 +1117,7 @@ where
 
     #[inline]
     fn prune(&mut self) {
-        self.branch.data.prune()
+        self.base.prune();
     }
 }
 

manta-crypto/src/merkle_tree/leaf_map.rs

@@ -108,6 +108,9 @@ where
     /// Returns a vector with all [`LeafDigest`]s, consuming `self`.
     fn into_leaf_digests(self) -> Vec<LeafDigest<C>>;
 
+    /// Returns a vector with all [`LeafDigest`]s and their markings, consuming `self`.
+    fn into_leaf_digests_with_markings(self) -> Vec<(bool, LeafDigest<C>)>;
+
     /// Returns a vector with all marked [`LeafDigest`]s in `self`.
     #[inline]
     fn marked_leaf_digests(&self) -> Vec<&LeafDigest<C>> {
@@ -138,7 +141,7 @@ where
 #[derivative(
     Clone(bound = "LeafDigest<C>: Clone"),
     Debug(bound = "LeafDigest<C>: Debug"),
-    Default(bound = "LeafDigest<C>: Default"),
+    Default(bound = ""),
     Eq(bound = "LeafDigest<C>: Eq"),
     Hash(bound = "LeafDigest<C>: Hash"),
     PartialEq(bound = "LeafDigest<C>: PartialEq")
@@ -178,7 +181,7 @@ where
 
     #[inline]
     fn push(&mut self, leaf_digest: LeafDigest<C>) {
-        self.0.push(leaf_digest);
+        self.0.push(leaf_digest)
     }
 
     #[inline]
@@ -196,6 +199,11 @@ where
         self.0
     }
 
+    #[inline]
+    fn into_leaf_digests_with_markings(self) -> Vec<(bool, LeafDigest<C>)> {
+        core::iter::repeat(false).zip(self.0).collect()
+    }
+
     #[inline]
     fn mark(&mut self, index: usize) {
         let _ = index;
@@ -230,7 +238,7 @@ where
 #[derivative(
     Clone(bound = "LeafDigest<C>: Clone"),
     Debug(bound = "LeafDigest<C>: Debug"),
-    Default(bound = "LeafDigest<C>: Default"),
+    Default(bound = ""),
     Eq(bound = "LeafDigest<C>: Eq"),
     Hash(bound = "LeafDigest<C>: Hash"),
     PartialEq(bound = "LeafDigest<C>: PartialEq")
@@ -274,7 +282,7 @@ where
 
     #[inline]
     fn push(&mut self, leaf_digest: LeafDigest<C>) {
-        self.vec.push((false, leaf_digest));
+        self.vec.push((false, leaf_digest))
     }
 
     #[inline]
@@ -295,6 +303,11 @@ where
         self.vec.into_iter().map(|(_, digest)| digest).collect()
     }
 
+    #[inline]
+    fn into_leaf_digests_with_markings(self) -> Vec<(bool, LeafDigest<C>)> {
+        self.vec
+    }
+
     #[inline]
     fn mark(&mut self, index: usize) {
         if let Some((b, _)) = self.vec.get_mut(index) {
@@ -331,8 +344,9 @@ where
 #[derivative(
     Clone(bound = "LeafDigest<C>: Clone"),
     Debug(bound = "LeafDigest<C>: Debug"),
-    Default(bound = "LeafDigest<C>: Default"),
+    Default(bound = ""),
     Eq(bound = "LeafDigest<C>: Eq"),
+    Hash(bound = "LeafDigest<C>: Hash"),
     PartialEq(bound = "LeafDigest<C>: PartialEq")
 )]
 pub struct LeafBTreeMap<C>
@@ -368,7 +382,10 @@ where
 
     #[inline]
     fn position(&self, leaf_digest: &LeafDigest<C>) -> Option<usize> {
-        self.map.iter().position(|(_, (_, l))| l == leaf_digest)
+        self.map
+            .keys()
+            .find(|index| self.get(**index) == Some(leaf_digest))
+            .copied()
     }
 
     #[inline]
@@ -409,6 +426,14 @@ where
             .collect()
     }
 
+    #[inline]
+    fn into_leaf_digests_with_markings(self) -> Vec<(bool, LeafDigest<C>)> {
+        self.map
+            .into_iter()
+            .map(|(_, (marking, digest))| (marking, digest))
+            .collect()
+    }
+
     #[inline]
     fn mark(&mut self, index: usize) {
         if let Some((b, _)) = self.map.get_mut(&index) {
@@ -425,7 +450,7 @@ where
     fn remove(&mut self, index: usize) -> bool {
         match self.last_index {
             Some(current_index) if index == current_index => false,
-            _ => !matches!(self.map.remove(&index), None),
+            _ => self.map.remove(&index).is_some(),
         }
     }
 }
@@ -449,7 +474,7 @@ where
 #[derivative(
     Clone(bound = "LeafDigest<C>: Clone"),
     Debug(bound = "LeafDigest<C>: Debug"),
-    Default(bound = "LeafDigest<C>: Default"),
+    Default(bound = ""),
     Eq(bound = "LeafDigest<C>: Eq"),
     PartialEq(bound = "LeafDigest<C>: PartialEq")
 )]
@@ -488,7 +513,10 @@ where
 
     #[inline]
     fn position(&self, leaf_digest: &LeafDigest<C>) -> Option<usize> {
-        self.map.iter().position(|(_, (_, l))| l == leaf_digest)
+        self.map
+            .keys()
+            .find(|index| self.get(**index) == Some(leaf_digest))
+            .copied()
     }
 
     #[inline]
@@ -529,6 +557,14 @@ where
             .collect()
     }
 
+    #[inline]
+    fn into_leaf_digests_with_markings(self) -> Vec<(bool, LeafDigest<C>)> {
+        self.map
+            .into_iter()
+            .map(|(_, (marking, digest))| (marking, digest))
+            .collect()
+    }
+
     #[inline]
     fn mark(&mut self, index: usize) {
         if let Some((b, _)) = self.map.get_mut(&index) {
@@ -545,7 +581,7 @@ where
     fn remove(&mut self, index: usize) -> bool {
         match self.last_index {
             Some(current_index) if index == current_index => false,
-            _ => !matches!(self.map.remove(&index), None),
+            _ => self.map.remove(&index).is_some(),
         }
     }
 }