factor out subtree-append logic

src/merkle/smt/mod.rs

@@ -89,7 +89,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         };
 
         Self::path_and_leaf_to_opening(merkle_path, leaf)
-}
+    }
 
     /// Inserts a value at the specified key, returning the previous value associated with that key.
     /// Recall that by definition, any key that hasn't been updated is associated with
@@ -114,7 +114,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
 
         old_value
-}
+    }
 
     /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
     /// `node_hash_at_index` is the hash of the node stored at index.
@@ -122,7 +122,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         &mut self,
         mut index: NodeIndex,
         node_hash_at_index: RpoDigest,
-) {
+    ) {
         let mut node_hash = node_hash_at_index;
         for node_depth in (0..index.depth()).rev() {
             let is_right = index.is_value_odd();
@@ -144,7 +144,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             }
         }
         self.set_root(node_hash);
-}
+    }
 
     /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
     /// tree, allowing for validation before applying those changes.
@@ -157,7 +157,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     fn compute_mutations(
         &self,
         kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
-) -> MutationSet<DEPTH, Self::Key, Self::Value> {
+    ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
         use NodeMutation::*;
 
         let mut new_root = self.root();
@@ -242,7 +242,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             node_mutations,
             new_pairs,
         }
-}
+    }
 
     /// Apply the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
     /// this tree.
@@ -255,10 +255,10 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     fn apply_mutations(
         &mut self,
         mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
-) -> Result<(), MerkleError>
+    ) -> Result<(), MerkleError>
-where
+    where
         Self: Sized,
-{
+    {
         use NodeMutation::*;
         let MutationSet {
             old_root,
@@ -287,7 +287,7 @@ where
         self.set_root(new_root);
 
         Ok(())
-}
+    }
 
     // REQUIRED METHODS
     // ---------------------------------------------------------------------------------------------
@@ -336,7 +336,7 @@ where
         existing_leaf: Self::Leaf,
         key: &Self::Key,
         value: &Self::Value,
-) -> Self::Leaf;
+    ) -> Self::Leaf;
 
     /// Maps a key to a leaf index
     fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
@@ -383,7 +383,7 @@ where
             let hash = Self::hash_leaf(&leaf);
 
             accumulator.nodes.insert(col, leaf);
-            accumulator.add_leaf(SubtreeLeaf { col, hash });
+            add_subtree_leaf(&mut accumulator.leaves, SubtreeLeaf { col, hash });
 
             debug_assert!(current_leaf_buffer.is_empty());
         }
@@ -631,6 +631,7 @@ impl SubtreeLeaf {
     }
 }
 
+/// Helper struct to organize the return value of [`SparseMerkleTree::sorted_pairs_to_leaves()`].
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct PairComputations<K, L> {
     /// Literal leaves to be added to the sparse Merkle tree's internal mapping.
@@ -679,6 +680,36 @@ impl<K, L> Default for PairComputations<K, L> {
     }
 }
 
+/// Handles the logic for figuring out whether the new leaf starts a new subtree or not.
+fn add_subtree_leaf(subtrees: &mut Vec<Vec<SubtreeLeaf>>, leaf: SubtreeLeaf) {
+    let last_subtree = match subtrees.last_mut() {
+        // Base case.
+        None => {
+            subtrees.push(vec![leaf]);
+            return;
+        },
+        Some(last_subtree) => last_subtree,
+    };
+
+    debug_assert!(!last_subtree.is_empty());
+    debug_assert!(last_subtree.len() <= COLS_PER_SUBTREE as usize);
+
+    // The multiple of 256 after 0 is 1, but 0 and 1 do not belong to different subtrees.
+    let last_subtree_col = u64::max(1, last_subtree.last().unwrap().col);
+    let next_subtree_col = if Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE) {
+        u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
+    } else {
+        last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
+    };
+
+    if leaf.col < next_subtree_col {
+        last_subtree.push(leaf);
+    } else {
+        let next_subtree = vec![leaf];
+        subtrees.push(next_subtree);
+    }
+}
+
 // TESTS
 // ================================================================================================
 #[cfg(test)]
@@ -687,9 +718,7 @@ mod test {
 
     use alloc::{collections::BTreeMap, vec::Vec};
 
-    use num::Integer;
+    use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf};
-
-    use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE};
     use crate::{
         hash::rpo::RpoDigest,
         merkle::{NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
@@ -904,25 +933,7 @@ mod test {
                 accumulated_nodes.extend(nodes);
 
                 for subtree_leaf in next_leaves {
-                    if leaf_subtrees.is_empty() {
+                    super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
-                        leaf_subtrees.push(vec![subtree_leaf]);
-                        continue;
-                    }
-
-                    let buffer_max_col =
-                        u64::max(1, leaf_subtrees.last().unwrap().last().unwrap().col);
-                    let next_subtree_col =
-                        if Integer::is_multiple_of(&buffer_max_col, &COLS_PER_SUBTREE) {
-                            u64::next_multiple_of(buffer_max_col + 1, COLS_PER_SUBTREE)
-                        } else {
-                            buffer_max_col.next_multiple_of(COLS_PER_SUBTREE)
-                        };
-
-                    if subtree_leaf.col < next_subtree_col {
-                        leaf_subtrees.last_mut().unwrap().push(subtree_leaf);
-                    } else {
-                        leaf_subtrees.push(vec![subtree_leaf]);
-                    }
                 }
             }