From 7ff9ab9314d42fb0061ed71789471fd3ee3355b9 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Mon, 4 Nov 2024 12:53:27 -0700
Subject: [PATCH] factor out subtree-append logic

---
 src/merkle/smt/mod.rs | 87 +++++++++++++++++--------------------------
 1 file changed, 34 insertions(+), 53 deletions(-)

diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index af39b08..1f857cf 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -383,7 +383,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             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.
@@ -639,36 +640,6 @@ pub(crate) struct PairComputations<K, L> {
     pub leaves: Vec<Vec<SubtreeLeaf>>,
 }
 
-impl<K, L> PrecomputedSubtrees<K, L> {
-    pub fn add_leaf(&mut self, leaf: SubtreeLeaf) {
-        let last_subtree = match self.leaves.last_mut() {
-            // Base case.
-            None => {
-                self.leaves.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];
-            self.leaves.push(next_subtree);
-        }
-    }
-}
-
 // Derive requires `L` to impl Default, even though we don't actually need that.
 impl<K, L> Default for PairComputations<K, L> {
     fn default() -> Self {
@@ -679,6 +650,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 +688,7 @@ mod test {
 
     use alloc::{collections::BTreeMap, vec::Vec};
 
-    use num::Integer;
-
-    use super::{SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE};
+    use super::{SparseMerkleTree, SubtreeLeaf};
     use crate::{
         hash::rpo::RpoDigest,
         merkle::{smt::InnerNode, NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
@@ -904,25 +903,7 @@ mod test {
                 accumulated_nodes.extend(nodes);
 
                 for subtree_leaf in next_leaves {
-                    if leaf_subtrees.is_empty() {
-                        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]);
-                    }
+                    super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
                 }
             }