From 2fac7f11d9ed5c110b03588eceb4a17131fa2a75 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Fri, 16 Aug 2024 19:22:15 -0600
Subject: [PATCH 1/2] WIP: slightly refactor
 recompute_nodes_from_index_to_root()

---
 src/merkle/smt/mod.rs | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index c6699ef..e6275cc 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -116,27 +116,24 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         mut index: NodeIndex,
         node_hash_at_index: RpoDigest,
     ) {
-        let mut node_hash = node_hash_at_index;
+        let mut new_child_hash = node_hash_at_index;
         for node_depth in (0..index.depth()).rev() {
             let is_right = index.is_value_odd();
             index.move_up();
-            let InnerNode { left, right } = self.get_inner_node(index);
-            let (left, right) = if is_right {
-                (left, node_hash)
-            } else {
-                (node_hash, right)
-            };
-            node_hash = Rpo256::merge(&[left, right]);
+            let old_node = self.get_inner_node(index);
+            let new_node = old_node.create_update(is_right, new_child_hash);
+            new_child_hash = new_node.hash();
 
-            if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
+            let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
+            if new_child_hash == equivalent_empty_hash {
                 // If a subtree is empty, when can remove the inner node, since it's equal to the
                 // default value
                 self.remove_inner_node(index)
             } else {
-                self.insert_inner_node(index, InnerNode { left, right });
+                self.insert_inner_node(index, new_node);
             }
         }
-        self.set_root(node_hash);
+        self.set_root(new_child_hash);
     }
 
     // REQUIRED METHODS
@@ -195,6 +192,17 @@ pub(crate) struct InnerNode {
 }
 
 impl InnerNode {
+    /// Creates a new node for when a child node has been changed.
+    pub fn create_update(self, update_right: bool, new_hash: RpoDigest) -> Self {
+        let (left, right) = if update_right {
+            (self.left, new_hash)
+        } else {
+            (new_hash, self.right)
+        };
+
+        Self { left, right }
+    }
+
     pub fn hash(&self) -> RpoDigest {
         Rpo256::merge(&[self.left, self.right])
     }

From 3d9c82bbe54aae53f156bf4f1bf046cbde2c230c Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Fri, 16 Aug 2024 19:31:48 -0600
Subject: [PATCH 2/2] add get_value() to SparseMerkleTree trait

---
 src/merkle/smt/full/mod.rs   | 9 +++++++++
 src/merkle/smt/mod.rs        | 3 +++
 src/merkle/smt/simple/mod.rs | 7 +++++++
 3 files changed, 19 insertions(+)

diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 2de5786..e5a60d6 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/src/merkle/smt/full/mod.rs
@@ -249,6 +249,15 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         }
     }
 
+    fn get_value(&self, key: &RpoDigest) -> Word {
+        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
+
+        self.leaves
+            .get(&leaf_pos)
+            .map(|leaf| leaf.get_value(key).unwrap_or_default())
+            .unwrap_or_default()
+    }
+
     fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
         let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
 
diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index e6275cc..8278667 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -157,6 +157,9 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// Inserts a leaf node, and returns the value at the key if already exists
     fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
 
+    /// Returns the value at the specified key, without modifying the tree.
+    fn get_value(&self, key: &Self::Key) -> Self::Value;
+
     /// Returns the leaf at the specified index.
     fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
 
diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs
index 9f0fed1..a4ff136 100644
--- a/src/merkle/smt/simple/mod.rs
+++ b/src/merkle/smt/simple/mod.rs
@@ -288,6 +288,13 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         }
     }
 
+    fn get_value(&self, key: &LeafIndex<DEPTH>) -> Word {
+        let leaf_pos = key.value();
+
+        *self.leaves.get(&leaf_pos)
+            .unwrap_or(&Self::EMPTY_VALUE)
+    }
+
     fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
         let leaf_pos = key.value();
         match self.leaves.get(&leaf_pos) {