sdk/crypto/smt: Move tree recalculation into separate internal function

src/sdk/src/crypto/smt/mod.rs

@@ -189,31 +189,7 @@ impl<
             dirty_idxs.push(parent_idx);
         }
 
-        // Depth first from the bottom of the tree
-        for _ in 0..N + 1 {
-            let mut new_dirty_idxs = Vec::new();
-
-            for idx in dirty_idxs {
-                let left_idx = util::left_child(&idx);
-                let right_idx = util::right_child(&idx);
-                let left = self.get_node(&left_idx);
-                let right = self.get_node(&right_idx);
-                // Recalculate the node
-                let node = self.hasher.hash([left, right]);
-
-                self.put_node(idx.clone(), node)?;
-
-                // Add this node's parent to the update list
-                let parent_idx = match util::parent(&idx) {
-                    Some(idx) => idx,
-                    // We are at the root node so no parents exist
-                    None => break,
-                };
-                new_dirty_idxs.push(parent_idx);
-            }
-
-            dirty_idxs = new_dirty_idxs;
-        }
+        self.recompute_tree(&mut dirty_idxs)?;
 
         Ok(())
     }
@@ -233,40 +209,46 @@ impl<
             dirty_idxs.push(parent_idx);
         }
 
-        // Depth first from the bottom of the tree
+        self.recompute_tree(&mut dirty_idxs)?;
+
+        Ok(())
+    }
+
+    /// Returns the Merkle tree root.
+    pub fn root(&self) -> F {
+        self.get_node(&BigUint::from(0u32))
+    }
+
+    /// Recomputes the Merkle tree depth first from the bottom of the tree
+    fn recompute_tree(&mut self, dirty_idxs: &mut Vec<BigUint>) -> ContractResult {
         for _ in 0..N + 1 {
-            let mut new_dirty_idxs = Vec::new();
+            let mut new_dirty_idxs = vec![];
 
-            for idx in dirty_idxs {
-                let left_idx = util::left_child(&idx);
-                let right_idx = util::right_child(&idx);
+            for idx in &mut *dirty_idxs {
+                let left_idx = util::left_child(idx);
+                let right_idx = util::right_child(idx);
                 let left = self.get_node(&left_idx);
                 let right = self.get_node(&right_idx);
-                // Recalculate the node
+                // Recalclate the node
                 let node = self.hasher.hash([left, right]);
-
                 self.put_node(idx.clone(), node)?;
 
                 // Add this node's parent to the update list
-                let parent_idx = match util::parent(&idx) {
+                let parent_idx = match util::parent(idx) {
                     Some(idx) => idx,
                     // We are at the root node so no parents exist
                     None => break,
                 };
+
                 new_dirty_idxs.push(parent_idx);
             }
 
-            dirty_idxs = new_dirty_idxs;
+            *dirty_idxs = new_dirty_idxs;
         }
 
         Ok(())
     }
 
-    /// Returns the Merkle tree root.
-    pub fn root(&self) -> F {
-        self.get_node(&BigUint::from(0u32))
-    }
-
     /// Give the path leading from the leaf at `index` up to the root. This is
     /// a "proof" in the sense of "valid path in a Merkle tree", not a ZK argument.
     pub fn prove_membership(&self, pos: &F) -> Path<N, F, H> {