MULTI INSERTION WORKING STATE

src/merkle/index.rs

@@ -147,6 +147,13 @@ impl NodeIndex {
         self.value >>= 1;
     }
 
+    pub fn as_up(&self) -> Self {
+        Self {
+            depth: self.depth.saturating_sub(1),
+            value: self.value >> 1,
+        }
+    }
+
     /// Traverses towards the root until the specified depth is reached.
     ///
     /// Assumes that the specified depth is smaller than the current depth.

src/merkle/smt/full/mod.rs

@@ -19,6 +19,8 @@ mod proof;
 pub use proof::SmtProof;
 use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
+mod rollback;
+
 #[cfg(test)]
 mod tests;
 
@@ -263,6 +265,11 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         }
     }
 
+    fn node_with_empty_children(&self, index: NodeIndex) -> InnerNode {
+        let &child = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1);
+        InnerNode { left: child, right: child }
+    }
+
     fn get_value(&self, key: &Self::Key) -> Self::Value {
         let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
 
@@ -281,6 +288,11 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         }
     }
 
+    fn maybe_get_leaf(&self, key: &RpoDigest) -> Option<Self::Leaf> {
+        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
+        self.leaves.get(&leaf_pos).cloned()
+    }
+
     fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
         leaf.hash()
     }
@@ -304,6 +316,27 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
         }
     }
 
+    fn get_prospective_leaf(
+        &self,
+        existing_leaf: Option<&SmtLeaf>,
+        key: &RpoDigest,
+        value: &Word,
+    ) -> SmtLeaf {
+        //let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(key);
+        match existing_leaf {
+            None => SmtLeaf::new_single(*key, *value),
+            Some(existing_leaf) => {
+                let mut new_leaf = existing_leaf.clone();
+                if *value != EMPTY_WORD {
+                    new_leaf.insert(*key, *value);
+                } else {
+                    new_leaf.remove(*key);
+                }
+                new_leaf
+            },
+        }
+    }
+
     fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
         let most_significant_felt = key[3];
         LeafIndex::new_max_depth(most_significant_felt.as_int())

src/merkle/smt/full/tests.rs

@@ -282,17 +282,26 @@ fn test_checked_insertion() {
 
     // insert key-value 1
     let root_1 = {
+        let prospective_leaf =
+            smt.get_prospective_leaf(smt.maybe_get_leaf(&key_1).as_ref(), &key_1, &value_1);
         let prospective = smt.hash_prospective_leaf(&key_1, &value_1);
         let old_value_1 = smt.insert(key_1, value_1);
         assert_eq!(old_value_1, EMPTY_WORD);
 
         assert_eq!(prospective, smt.get_leaf(&key_1).hash());
+        assert_eq!(prospective_leaf, smt.get_leaf(&key_1));
 
         assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
 
         smt.root()
     };
 
+    let parent_1_index = {
+        let node_index = Smt::key_to_leaf_index(&key_1);
+        NodeIndex::from(node_index).as_up()
+    };
+    let leaf_1_parent = smt.get_inner_node(parent_1_index);
+
     {
         // Trying to insert something else into key_1 with the existing root should fail, and
         // should not modify the tree at all.
@@ -313,11 +322,15 @@ fn test_checked_insertion() {
 
     // insert key-value 2
     let root_2 = {
+        let prospective_leaf =
+            smt.get_prospective_leaf(smt.maybe_get_leaf(&key_2).as_ref(), &key_2, &value_2);
         let prospective = smt.hash_prospective_leaf(&key_2, &value_2);
+        assert_eq!(prospective_leaf.hash(), prospective);
         let old_value_2 = smt.insert(key_2, value_2);
         assert_eq!(old_value_2, EMPTY_WORD);
 
         assert_eq!(prospective, smt.get_leaf(&key_2).hash());
+        assert_eq!(prospective_leaf, smt.get_leaf(&key_2));
 
         assert_eq!(
             smt.get_leaf(&key_2),
@@ -342,11 +355,15 @@ fn test_checked_insertion() {
 
     // insert key-value 3
     let root_3 = {
+        let prospective_leaf =
+            smt.get_prospective_leaf(smt.maybe_get_leaf(&key_3).as_ref(), &key_3, &value_3);
         let prospective = smt.hash_prospective_leaf(&key_3, &value_3);
+        assert_eq!(prospective_leaf.hash(), prospective);
         let old_value_3 = smt.insert(key_3, value_3);
         assert_eq!(old_value_3, EMPTY_WORD);
 
         assert_eq!(prospective, smt.get_leaf(&key_3).hash());
+        assert_eq!(prospective_leaf, smt.get_leaf(&key_3));
 
         assert_eq!(
             smt.get_leaf(&key_3),
@@ -371,6 +388,8 @@ fn test_checked_insertion() {
 
     // remove key 3
     {
+        let prospective_leaf =
+            smt.get_prospective_leaf(smt.maybe_get_leaf(&key_3).as_ref(), &key_3, &EMPTY_WORD);
         let old_hash = smt.get_leaf(&key_3).hash();
         let old_value_3 = smt.insert_ensure_root(key_3, EMPTY_WORD, root_2).unwrap();
         assert_eq!(old_value_3, value_3);
@@ -381,6 +400,8 @@ fn test_checked_insertion() {
             SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
         );
 
+        assert_eq!(prospective_leaf, smt.get_leaf(&key_3));
+
         assert_eq!(smt.root(), root_2);
     }
 
@@ -407,6 +428,78 @@ fn test_checked_insertion() {
 
         assert_eq!(smt.root(), root_empty);
     }
+
+    let mut mutations = smt.start_muts();
+    let smt = Smt::default();
+    assert_eq!(smt, smt_empty);
+
+    let mut control = smt.clone();
+    control.insert(key_1, value_1);
+    assert_eq!(control.root(), root_1);
+
+    smt.good_insert(&mut mutations, key_1, value_1);
+    assert_eq!(mutations.new_root, root_1);
+    let new_parent_1 = mutations.get_inner_node(parent_1_index).unwrap();
+    assert_eq!(leaf_1_parent, new_parent_1);
+    let mut index = NodeIndex::from(Smt::key_to_leaf_index(&key_1));
+    for depth in (0..smt.depth()).rev() {
+        let is_right = index.is_value_odd();
+        index.move_up();
+        assert_eq!(depth, index.depth());
+
+        let test = mutations.get_inner_node(index);
+
+        let &equivalent_empty_hash = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth());
+        let control_node = control.get_inner_node(index);
+
+        match test {
+            Some(ref existing) => assert_eq!(&control_node, existing),
+            None => {
+                assert_eq!(control_node, smt.node_with_empty_children(index));
+            },
+        }
+
+        //if control.hash() == equivalent_empty_hash {
+        if control_node == smt.node_with_empty_children(index) {
+            assert_eq!(control_node, smt.node_with_empty_children(index));
+            assert_eq!(control_node.hash(), equivalent_empty_hash);
+            match test {
+                Some(test) => assert_eq!(control_node, test, "at depth index {index:?}"),
+                None => assert_eq!(None, test, "at depth index {index:?}"),
+            }
+            //assert_eq!(None, test, "at depth index {index:?}");
+        } else {
+            assert_eq!(Some(control_node), test, "at depth index {index:?}");
+        }
+    }
+
+    control.insert(key_2, value_2);
+    assert_eq!(control.root(), root_2);
+
+    std::eprintln!("\nTESTING SECOND INSERT");
+    smt.good_insert(&mut mutations, key_2, value_2);
+
+    let parent_2_index = {
+        let node_index = Smt::key_to_leaf_index(&key_2);
+        NodeIndex::from(node_index).as_up()
+    };
+    assert_eq!(
+        control.get_inner_node(parent_1_index),
+        mutations.get_inner_node(parent_1_index).unwrap()
+    );
+    let leaf_2_parent = control.get_inner_node(parent_2_index);
+    let new_parent_2 = mutations.get_inner_node(parent_2_index).unwrap();
+    assert_eq!(leaf_2_parent, new_parent_2);
+    let mut index = NodeIndex::from(Smt::key_to_leaf_index(&key_2));
+    for depth in (0..smt.depth()).rev() {
+        index.move_up();
+        assert_eq!(depth, index.depth());
+
+        let test = mutations.get_inner_node(index);
+        let control_node = control.get_inner_node(index);
+    }
+
+    assert_eq!(mutations.new_root, root_2);
 }
 
 /// Tests that 2 key-value pairs stored in the same leaf have the same path

src/merkle/smt/mod.rs

@@ -1,4 +1,8 @@
-use alloc::vec::Vec;
+use alloc::{
+    collections::{btree_map::Entry, BTreeMap},
+    vec::Vec,
+};
+use tap::Pipe;
 
 use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
 use crate::{
@@ -12,6 +16,18 @@ pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
 mod simple;
 pub use simple::SimpleSmt;
 
+macro_rules! func {
+    () => {{
+        fn __f() {}
+        fn type_name_of<T>(_: T) -> &'static str {
+            std::any::type_name::<T>()
+        }
+        let name = type_name_of(__f);
+        let name = name.strip_suffix("::__f").unwrap();
+        format!("{}()", name.split("::").last().unwrap())
+    }};
+}
+
 // CONSTANTS
 // ================================================================================================
 
@@ -43,13 +59,13 @@ pub const SMT_MAX_DEPTH: u8 = 64;
 /// must accomodate all keys that map to the same leaf.
 ///
 /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
-pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
+pub(crate) trait SparseMerkleTree<const DEPTH: u8>: Clone {
     /// The type for a key
-    type Key: Clone;
+    type Key: Clone + PartialEq + alloc::fmt::Debug;
     /// The type for a value
     type Value: Clone + PartialEq;
     /// The type for a leaf
-    type Leaf;
+    type Leaf: alloc::fmt::Debug + PartialEq;
     /// The type for an opening (i.e. a "proof") of a leaf
     type Opening;
 
@@ -107,6 +123,13 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
 
         self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
 
+        if cfg!(debug_assertions) || cfg!(test) {
+            let is_right = node_index.is_value_odd();
+            let parent = self.get_inner_node(node_index.as_up());
+            let leaf_hash_from_node = if is_right { parent.right } else { parent.left };
+            assert_eq!(leaf_hash_from_node, Self::hash_leaf(&leaf));
+        }
+
         old_value
     }
 
@@ -131,21 +154,240 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             return Ok(value);
         }
 
-        // Compute the nodes we'll need to make and remove.
+        //let mut mutations = Mutations::new();
-        let mut removals: Vec<NodeIndex> = Vec::with_capacity(DEPTH as usize);
+        //self.prospective_insert(&mut mutations, key.clone(), value.clone());
-        let mut additions: Vec<(NodeIndex, InnerNode)> = Vec::with_capacity(DEPTH as usize);
+        //let Mutations {
+        //    mut node_removals,
+        //    mut node_additions,
+        //    mut pair_insertions,
+        //    new_root,
+        //} = mutations;
+
+        let mut control = self.clone();
+        control.insert(key.clone(), value.clone());
+
+        let mut index = NodeIndex::from(Self::key_to_leaf_index(&key));
+
+        let mut mutations = self.start_muts();
+        self.good_insert(&mut mutations, key, value);
+
+        for depth in (0..index.depth()).rev() {
+            index.move_up();
+            assert_eq!(depth, index.depth());
+
+            let control_node = control.get_inner_node(index);
+            let test_node = mutations.maybe_get_node(index).or_empty();
+            //let test_node = match mutations.get_inner_node(index) {
+            //    Some(node) => node,
+            //    None => {
+            //        let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+            //        InnerNode { left: child, right: child }
+            //    },
+            //};
+
+            assert_eq!(control_node, test_node, "nodes at {index:?} are different!");
+        }
+
+        assert_eq!(control.root(), mutations.new_root);
+
+        if expected_root != mutations.new_root {
+            return Err(MerkleError::ConflictingRoots(vec![expected_root, mutations.new_root]));
+        }
+
+        self.apply_muts(mutations);
+
+        // Actual mutations start here.
+        //self.apply_mutations(mutations);
+
+        //for (key, value) in pair_insertions.drain(..) {
+        //    self.insert_value(key, value);
+        //}
+        //
+        //for index in node_removals.drain(..) {
+        //    self.remove_inner_node(index);
+        //}
+        //
+        //for (index, new_node) in node_additions.drain(..) {
+        //    self.insert_inner_node(index, new_node);
+        //}
+        //
+        //self.set_root(new_root);
+
+        Ok(old_value)
+    }
+
+    //fn prospective_insert(
+    //    &self,
+    //    mutations: &mut Mutations<DEPTH, Self::Key, Self::Value>,
+    //    key: Self::Key,
+    //    value: Self::Value,
+    //) {
+    //    let old_value = self.get_value(&key);
+    //    // If the old value and new value are the same, there's nothing to update.
+    //    if value == old_value {
+    //        return;
+    //    }
+    //
+    //    let mut node_index = {
+    //        let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
+    //        NodeIndex::from(leaf_index)
+    //    };
+    //
+    //    //let mut new_child_hash = self.hash_prospective_leaf(&key, &value);
+    //    //let existing_leaf = match self.maybe_get_leaf(&key) {
+    //    //    Some(leaf) => Some(leaf),
+    //    //    None => mutations
+    //    //        .pair_insertions
+    //    //        .iter()
+    //    //        .find(|&pair| pair.0 == key)
+    //    //        .cloned()
+    //    //        .map(|pair| self.get_prospective_leaf(None, &pair.0, &pair.1)),
+    //    //};
+    //    let maybe_prospective_leaf = mutations
+    //        .pair_insertions
+    //        .iter()
+    //        .find(|&pair| pair.0 == key)
+    //        .cloned()
+    //        .map(|pair| self.get_prospective_leaf(None, &pair.0, &pair.1));
+    //    std::dbg!(&maybe_prospective_leaf);
+    //    let existing_leaf = match maybe_prospective_leaf {
+    //        Some(leaf) => Some(leaf),
+    //        None => self.maybe_get_leaf(&key),
+    //    };
+    //    let new_child = self.get_prospective_leaf(existing_leaf.as_ref(), &key, &value);
+    //    let mut new_child_hash = Self::hash_leaf(&new_child);
+    //    for node_depth in (0..node_index.depth()).rev() {
+    //        let is_right = node_index.is_value_odd();
+    //        node_index.move_up();
+    //
+    //        //let old_node = self.get_inner_node(node_index);
+    //        //let old_node = mutations.get_inner_node(node_index).unwrap_or_else(|| {
+    //        //    std::eprintln!("nothing in mutations, getting existing");
+    //        //    self.get_inner_node(node_index)
+    //        //});
+    //        let old_node = match mutations.get_inner_node(node_index) {
+    //            Some(old_node) => {
+    //                std::eprintln!("old node found in mutations, using that");
+    //                old_node
+    //            },
+    //            None => self.get_inner_node(node_index),
+    //        };
+    //        let new_node = if is_right {
+    //            InnerNode {
+    //                left: old_node.left,
+    //                right: new_child_hash,
+    //            }
+    //        } else {
+    //            InnerNode {
+    //                left: new_child_hash,
+    //                right: old_node.right,
+    //            }
+    //        };
+    //        if let Some(prospective_node) = mutations.get_inner_node_mut(node_index) {
+    //            *prospective_node = new_node.clone();
+    //        }
+    //
+    //        // The next iteration will operate on this node's new hash.
+    //        new_child_hash = new_node.hash();
+    //
+    //        let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
+    //        if new_child_hash == equivalent_empty_hash {
+    //            // If a subtree is empty, we can remove the inner node, since it's equal to the
+    //            // default value.
+    //            mutations.node_removals.push(node_index);
+    //        } else {
+    //            mutations.node_additions.push((node_index, new_node));
+    //        }
+    //    }
+    //
+    //    // Once we're at depth 0, the last node we made is the new root.
+    //    mutations.new_root = new_child_hash;
+    //    mutations.pair_insertions.push((key, value));
+    //}
+
+    fn good_insert(
+        &self,
+        mutations: &mut Muts<DEPTH, Self::Key, Self::Value>,
+        key: Self::Key,
+        value: Self::Value,
+    ) {
+        let old_value = self.get_value(&key);
+        if value == old_value {
+            return;
+        }
+
+        std::eprintln!("{}: performing insert with key {key:?}", func!());
 
         let mut node_index = {
             let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
             NodeIndex::from(leaf_index)
         };
 
-        let mut new_child_hash = self.hash_prospective_leaf(&key, &value);
+        let mut control = self.clone();
+        control.apply_muts(mutations.clone());
+        //control.insert(key.clone(), value.clone());
+
+        //let existing_leaf = self.maybe_get_leaf(&key);
+        //let mut maybe_leaf = mutations.maybe_get_values(&key);
+        //assert!(matches!(maybe_leaf.len(), 0 | 1));
+        //let maybe_leaf = maybe_leaf.pop();
+        //let existing_leaf = match maybe_leaf {
+        //    Some(existing_value) => {
+        //        let prospective_leaf =
+        //            self.get_prospective_leaf(existing_leaf.as_ref(), &key, &existing_value);
+        //        assert_eq!(existing_leaf, None);
+        //        Some(prospective_leaf)
+        //    },
+        //    None => existing_leaf,
+        //};
+        //
+        //let new_child = self.get_prospective_leaf(existing_leaf.as_ref(), &key, &value);
+        //assert_eq!(self.maybe_get_leaf(&key), None, "key {key:?} already has values?");
+        //let new_leaf = self.get_prospective_leaf(None, &key, &value);
+        let old_leaf = mutations.maybe_get_leaf(self, &key);
+        assert_eq!(control.maybe_get_leaf(&key), old_leaf);
+        if mutations.pair_insertions.is_empty() {
+            std::eprintln!(
+                "no insertions yet, our current leaf is {:?}",
+                self.maybe_get_leaf(&key)
+            );
+        }
+        let new_leaf = self.get_prospective_leaf(old_leaf.as_ref(), &key, &value);
+        let leaf_hash = Self::hash_leaf(&new_leaf);
+        assert_eq!(control.hash_prospective_leaf(&key, &value), leaf_hash);
+        let mut new_child_hash = leaf_hash;
+
+        // Prospective values for this key.
+        //let prospective_values = mutations.maybe_get_values(&key);
+        //assert_eq!(prospective_values.len(), 0);
+        //let new_leaf = match mutations.maybe_get_values(&key) {
+        //
+        //}
+
+        //let ref control_leaf = control.get_leaf(&key);
+        //assert_eq!(control_leaf, &existing_leaf.unwrap());
+        //let control_hash = Self::hash_leaf(control_leaf);
+        //assert_eq!(control_hash, new_child_hash);
+        //
+        //assert_eq!(control_leaf, &new_child);
+
         for node_depth in (0..node_index.depth()).rev() {
             let is_right = node_index.is_value_odd();
             node_index.move_up();
 
-            let old_node = self.get_inner_node(node_index);
+            let old_node = match mutations.get_inner_node(node_index) {
+                Some(prospective) => {
+                    //let existing = self.get_inner_node(node_index);
+                    //let &equivalent_empty = EmptySubtreeRoots::entry(DEPTH, node_depth);
+                    //if existing.hash() != equivalent_empty && existing != prospective {
+                    //    //assert_eq!(existing, prospective);
+                    //}
+                    prospective
+                },
+                None => self.get_inner_node(node_index),
+            };
+            assert_eq!(control.get_inner_node(node_index), old_node);
+
             let new_node = if is_right {
                 InnerNode {
                     left: old_node.left,
@@ -158,41 +400,134 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
                 }
             };
 
+            if cfg!(debug_assertions) || cfg!(test) {
+                let mut control = self.clone();
+                let mutations = mutations.clone();
+                control.apply_muts(mutations);
+                let control_old_node = control.get_inner_node(node_index);
+                assert_eq!(control_old_node, old_node);
+
+                //control.insert(key.clone(), value.clone());
+                //let control_leaf = control.get_leaf(&key);
+                //assert_eq!(control_leaf, new_leaf);
+                //let control_new_node = control.get_inner_node(node_index);
+                //assert_eq!(control_new_node, new_node);
+            }
+
+            let mut modified_control = control.clone();
+            modified_control.insert(key.clone(), value.clone());
+            let control_node = modified_control.get_inner_node(node_index);
+            assert_eq!(control_node, new_node, "is_right: {is_right:?}, old_node: {old_node:?}");
+
             // The next iteration will operate on this node's new hash.
             new_child_hash = new_node.hash();
+            std::eprintln!("\n{}: new_child_hash={new_child_hash:?}", func!());
 
             let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
             if new_child_hash == equivalent_empty_hash {
-                // If a subtree is empty, we can remove the inner node, since it's equal to the
+                mutations.maybe_remove_node(node_index);
-                // default value.
+
-                removals.push(node_index);
+                //mutations.mutations.push(Mutation::Removal(node_index));
+                //std::eprintln!("did removal do something? {did_change}");
             } else {
-                additions.push((node_index, new_node));
+                mutations.maybe_insert_inner_node(node_index, new_node);
+
+                let leaf_index = NodeIndex::from(Self::key_to_leaf_index(&key));
+                let leaf_parent = mutations.get_inner_node(leaf_index.as_up()).unwrap();
+                let hash_from_node = if leaf_index.is_value_odd() {
+                    leaf_parent.right
+                } else {
+                    leaf_parent.left
+                };
+                assert_eq!(hash_from_node, leaf_hash);
             }
         }
 
+        control.insert(key.clone(), value.clone());
+
         // Once we're at depth 0, the last node we made is the new root.
-        let new_root = new_child_hash;
+        mutations.new_root = new_child_hash;
+        mutations.pair_insertions.push((key, value));
 
-        if expected_root != new_root {
+        assert_eq!(control.root(), mutations.new_root);
-            return Err(MerkleError::ConflictingRoots(vec![expected_root, new_root]));
     }
 
-        // Actual mutations start here.
+    fn start_muts(&self) -> Muts<DEPTH, Self::Key, Self::Value> {
+        Muts::<DEPTH, Self::Key, Self::Value> {
+            //mutations: Default::default(),
+            node_mutations: Default::default(),
+            pair_insertions: Default::default(),
+            new_root: self.root(),
+        }
+    }
 
+    //fn apply_mutations(&mut self, mut mutations: Mutations<DEPTH, Self::Key, Self::Value>) {
+    //    for (key, value) in mutations.pair_insertions.drain(..) {
+    //        self.insert_value(key, value);
+    //    }
+    //
+    //    for index in mutations.node_removals.drain(..) {
+    //        self.remove_inner_node(index);
+    //    }
+    //
+    //    for (index, new_node) in mutations.node_additions.drain(..) {
+    //        self.insert_inner_node(index, new_node);
+    //    }
+    //
+    //    self.set_root(mutations.new_root);
+    //}
+
+    fn apply_muts(&mut self, mutations: Muts<DEPTH, Self::Key, Self::Value>) {
+        use NodeMutation::*;
+        let Muts {
+            node_mutations,
+            pair_insertions,
+            new_root,
+        } = mutations;
+
+        #[cfg(any(debug_assertions, test))]
+        let mut control = self.clone();
+        if cfg!(any(debug_assertions, test)) {
+            for (key, value) in pair_insertions.iter().cloned() {
+                control.insert(key, value);
+            }
+        }
+
+        for (index, mutation) in node_mutations.into_iter() {
+            match mutation {
+                Removal => self.remove_inner_node(index),
+                Addition(node) => self.insert_inner_node(index, node),
+            }
+            //if cfg!(any(debug_assertions, test)) {
+            //    let control_node = control.get_inner_node(index);
+            //    let test_node = self.get_inner_node(index);
+            //    assert_eq!(control_node, test_node);
+            //}
+        }
+        for (key, value) in pair_insertions.into_iter() {
             self.insert_value(key, value);
-
-        for index in removals.drain(..) {
-            self.remove_inner_node(index);
         }
-
-        for (index, new_node) in additions.drain(..) {
-            self.insert_inner_node(index, new_node);
-        }
-
         self.set_root(new_root);
 
-        Ok(old_value)
+        if cfg!(any(debug_assertions, test)) {
+            assert_eq!(control.root(), self.root());
+        }
+        //for mutation in mutations.node_mutations.drain(..) {
+        //    match mutation {
+        //        Removal(idx) => self.remove_inner_node(idx),
+        //        Addition(idx, node)
+        //    }
+        //}
+        //use Mutation::*;
+        //for mutation in mutations.mutations.drain(..) {
+        //    match mutation {
+        //        Removal(idx) => self.remove_inner_node(idx),
+        //        Addition(idx, node) => self.insert_inner_node(idx, node),
+        //        PairInsertion(key, value) => {
+        //            self.insert_value(key, value);
+        //        },
+        //    }
+        //}
     }
 
     /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
@@ -246,6 +581,8 @@ 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>;
 
+    fn node_with_empty_children(&self, index: NodeIndex) -> InnerNode;
+
     /// Returns the value at the specified key. Recall that by definition, any key that hasn't been
     /// updated is associated with [`Self::EMPTY_VALUE`].
     fn get_value(&self, key: &Self::Key) -> Self::Value;
@@ -253,6 +590,8 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// Returns the leaf at the specified index.
     fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
 
+    fn maybe_get_leaf(&self, key: &Self::Key) -> Option<Self::Leaf>;
+
     /// Returns the hash of a leaf
     fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
 
@@ -264,6 +603,13 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     /// the leaf each time.
     fn hash_prospective_leaf(&self, key: &Self::Key, value: &Self::Value) -> RpoDigest;
 
+    fn get_prospective_leaf(
+        &self,
+        existing_leaf: Option<&Self::Leaf>,
+        key: &Self::Key,
+        value: &Self::Value,
+    ) -> Self::Leaf;
+
     /// Maps a key to a leaf index
     fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
 
@@ -341,3 +687,340 @@ impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
         Self::new(node_index.value())
     }
 }
+
+// MUTATIONS
+//#[derive(Debug, Clone, PartialEq)]
+//#[cfg(not(test))]
+//pub enum Mutation<K, V> {
+//    Removal(NodeIndex),
+//    Addition(NodeIndex, InnerNode),
+//    PairInsertion(K, V),
+//}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum NodeMutation {
+    Removal,
+    Addition(InnerNode),
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Muts<const DEPTH: u8, K, V> {
+    //mutations: Vec<Mutation<K, V>>,
+    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
+    pair_insertions: Vec<(K, V)>,
+    new_root: RpoDigest,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct MaybeNode<const DEPTH: u8> {
+    inner: Option<InnerNode>,
+    index: NodeIndex,
+}
+impl<const DEPTH: u8> MaybeNode<DEPTH> {
+    pub fn from(inner: InnerNode, index: NodeIndex) -> Self {
+        Self { inner: Some(inner), index }
+    }
+
+    pub fn new_empty(index: NodeIndex) -> Self {
+        Self { inner: None, index }
+    }
+
+    pub fn or_empty(self) -> InnerNode {
+        match self.inner {
+            Some(node) => node,
+            None => {
+                let &child = EmptySubtreeRoots::entry(DEPTH, self.index.depth() + 1);
+                InnerNode { left: child, right: child }
+            },
+        }
+    }
+}
+
+impl<const DEPTH: u8, K: PartialEq, V: Clone> Muts<DEPTH, K, V> {
+    fn maybe_get_values(&self, key: &K) -> Vec<V> {
+        self.pair_insertions
+            .iter()
+            .filter_map(|(cur_key, cur_value)| if cur_key == key { Some(cur_value) } else { None })
+            .cloned()
+            .collect()
+    }
+
+    fn maybe_get_leaf<T>(&self, tree: &T, key: &K) -> Option<T::Leaf>
+    where
+        T: SparseMerkleTree<DEPTH, Key = K, Value = V>,
+    {
+        let leaf_index = T::key_to_leaf_index(key);
+        let pairs_at_index = self
+            .pair_insertions
+            .iter()
+            .filter(|&pair| T::key_to_leaf_index(&pair.0) == leaf_index);
+
+        let inital_acc: Option<T::Leaf> = tree.maybe_get_leaf(&key);
+        let leaf = pairs_at_index.fold(inital_acc, |acc, (k, v)| {
+            //
+            Some(tree.get_prospective_leaf(acc.as_ref(), k, v))
+        });
+
+        leaf
+    }
+
+    fn get_inner_node(&self, index: NodeIndex) -> Option<InnerNode> {
+        use NodeMutation::*;
+        match self.node_mutations.get(&index) {
+            Some(Addition(node)) => Some(node.clone()),
+            Some(Removal) => {
+                //std::eprintln!("{}: found other mutation? wat do. {other:?}", func!());
+                //panic!();
+                //let &empty_child_hash = EmptySubtreeRoots::entry(DEPTH, index.depth());
+                let node = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+                Some(InnerNode { left: *node, right: *node })
+            },
+            None => None,
+        }
+
+        //use Mutation::*;
+        //self.mutations.iter().find_map(|mutation| match mutation {
+        //    Addition(idx, node) if *idx == index => Some(node.clone()),
+        //    _ => None,
+        //})
+    }
+
+    fn maybe_get_node(&self, index: NodeIndex) -> MaybeNode<DEPTH> {
+        use NodeMutation::*;
+        match self.node_mutations.get(&index) {
+            Some(Addition(node)) => MaybeNode::from(node.clone(), index),
+            Some(Removal) => {
+                MaybeNode::from(MaybeNode::<DEPTH> { inner: None, index }.or_empty(), index)
+            },
+            None => MaybeNode::new_empty(index),
+        }
+    }
+
+    fn get_inner_node_mut(&mut self, index: NodeIndex) -> Option<&mut InnerNode> {
+        //use Mutation::*;
+        use NodeMutation::*;
+        //self.mutations.iter_mut().find_map(|mutation| match mutation {
+        //    Addition(idx, node) if *idx == index => Some(node),
+        //    Removal(idx) if *idx == index => {
+        //        panic!();
+        //    },
+        //    _ => None,
+        //})
+        match self.node_mutations.get_mut(&index) {
+            Some(Addition(node)) => Some(node),
+            _ => None,
+        }
+    }
+
+    pub fn maybe_insert_inner_node(&mut self, index: NodeIndex, node: InnerNode) -> bool {
+        //use Mutation::*;
+        use NodeMutation::*;
+
+        std::eprintln!(
+            "attempting to insert inner node at {index:?}, with {} mutations",
+            self.node_mutations.len()
+        );
+
+        let mut mutated: bool = false;
+        //let mut found_existing: bool = false;
+
+        // XXX
+        let orig = node.clone();
+
+        let mut node = Some(node);
+
+        use Entry::*;
+        match self.node_mutations.entry(index) {
+            Vacant(entry) => {
+                let node = node.take().unwrap();
+                std::eprintln!("{}: normal insert", func!());
+                entry.insert(Addition(node));
+            },
+            Occupied(mut entry) => match entry.get_mut() {
+                Addition(ref existing) if Some(existing) == node.as_ref() => {
+                    std::eprintln!("{}: existing entry is identical; doing nothing", func!());
+                },
+                Addition(existing) => {
+                    let node = node.take().unwrap();
+                    std::eprintln!(
+                        "{}: modifying existing addition entry from {:?} to {:?}",
+                        func!(),
+                        existing.hash(),
+                        node.hash()
+                    );
+                    mutated = true;
+                    *existing = node;
+                },
+                Removal => {
+                    std::eprintln!("{}: removal found; changing to addition", func!());
+                    mutated = true;
+                    entry.insert(Addition(node.take().unwrap()));
+                },
+            },
+        };
+
+        std::eprintln!("");
+
+        //match self.node_mutations.get_mut(&index) {
+        //    Some(mutation) => match mutation {
+        //        Addition(existing) => {
+        //            *existing = node.take().unwrap();
+        //            found_existing = true;
+        //            mutated = true;
+        //        },
+        //        Removal(existing) => {
+        //            self.node_mutations.remove()
+        //        }
+        //    },
+        //};
+
+        // XXX: sanity check
+        match self.get_inner_node(index) {
+            Some(inserted) => assert_eq!(orig, inserted),
+            None => {
+                let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+                let node = InnerNode { left: child, right: child };
+                assert_eq!(orig, node);
+            },
+        };
+
+        mutated
+
+        //self.mutations.retain_mut(|mutation: &mut Mutation<_, _>| match mutation {
+        //    Removal(idx) if *idx == index => {
+        //        mutated = true;
+        //        std::eprintln!("maybe_insert_node(): found removal at same index, removing");
+        //        false
+        //    },
+        //    Addition(idx, cur_node) => {
+        //        // FIXME: unnecessary clones
+        //        if Some(cur_node.clone()) == node.as_mut().cloned() {
+        //            found_existing = true;
+        //            std::eprintln!("maybe_insert_node(): found identical addition, ignoring");
+        //            panic!();
+        //        } else if *idx == index {
+        //            mutated = true;
+        //            std::eprintln!(
+        //                "maybe_insert_node(): found different addition at same index, replacing"
+        //            );
+        //            *cur_node = node.take().unwrap(); // XXX
+        //        }
+        //        true
+        //    },
+        //    _ => true,
+        //});
+        //
+        //if !mutated && !found_existing {
+        //    mutated = true;
+        //    self.mutations.push(Addition(index, node.take().unwrap()));
+        //}
+        //mutated
+    }
+
+    pub fn maybe_remove_node(&mut self, index: NodeIndex) -> bool {
+        use NodeMutation::*;
+        std::eprintln!(
+            "attempting to remove inner node at {index:?} with {} mutations",
+            self.node_mutations.len()
+        );
+
+        let mut mutated: bool = false;
+
+        use Entry::*;
+        match self.node_mutations.entry(index) {
+            Vacant(entry) => {
+                std::eprintln!("{}: normal removal", func!());
+                mutated = true;
+                entry.insert(Removal);
+            },
+            Occupied(mut entry) => match entry.get_mut() {
+                Addition(_) => {
+                    std::eprintln!("{}: found addition; changing to removal", func!());
+                    mutated = true;
+                    entry.insert(Removal);
+                },
+                Removal => {
+                    std::eprintln!("{}: identical to existing removal; doing nothing", func!());
+                },
+            },
+        };
+
+        // XXX: sanity check
+        match self.get_inner_node(index) {
+            Some(inserted) => {
+                let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+                let node = InnerNode { left: child, right: child };
+
+                assert_eq!(inserted, node);
+            },
+            None => (),
+        };
+
+        mutated
+
+        //    use Mutation::*;
+        //
+        //    let mut mutated: bool = false;
+        //    let mut found_existing: bool = false;
+        //
+        //    self.mutations.retain_mut(|mutation: &mut Mutation<_, _>| match mutation {
+        //        Removal(idx) if *idx == index => {
+        //            found_existing = true;
+        //            std::eprintln!(
+        //                "maybe_remove_node(): found a removal with same index, doing nothing"
+        //            );
+        //            true
+        //        },
+        //        Addition(idx, _) if *idx == index => {
+        //            mutated = true;
+        //            std::eprintln!("maybe_remove_node(): found an addition with same index, removing.");
+        //            false
+        //        },
+        //        _ => true,
+        //    });
+        //
+        //    if !found_existing {
+        //        mutated = true;
+        //        std::eprintln!("maybe_remove_node(): pushing Removal");
+        //        self.mutations.push(Removal(index));
+        //    }
+        //    mutated
+    }
+}
+
+//#[derive(Debug, Clone, PartialEq)]
+//pub struct Mutations<const DEPTH: u8, K, V> {
+//    node_removals: Vec<NodeIndex>,
+//    node_additions: Vec<(NodeIndex, InnerNode)>,
+//    pair_insertions: Vec<(K, V)>,
+//    new_root: RpoDigest,
+//}
+//
+//impl<const DEPTH: u8, K, V> Mutations<DEPTH, K, V> {
+//    pub fn new() -> Self {
+//        Mutations {
+//            node_removals: vec![],
+//            node_additions: vec![],
+//            pair_insertions: vec![],
+//            new_root: Default::default(),
+//        }
+//    }
+//
+//    pub fn get_inner_node(&self, index: NodeIndex) -> Option<InnerNode> {
+//        self.node_additions
+//            .iter()
+//            .find(|(idx, _)| *idx == index)
+//            .map(|(_, node)| node.clone())
+//        //.unwrap_or_else(|| {
+//        //    let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+//        //    InnerNode { left: child, right: child }
+//        //})
+//    }
+//
+//    pub fn get_inner_node_mut(&mut self, index: NodeIndex) -> Option<&mut InnerNode> {
+//        self.node_additions
+//            .iter_mut()
+//            .find(|(idx, _)| *idx == index)
+//            .map(|(_, node)| node)
+//    }
+//}

src/merkle/smt/simple/mod.rs

@@ -299,6 +299,11 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         let _ = self.inner_nodes.remove(&index);
     }
 
+    fn node_with_empty_children(&self, index: NodeIndex) -> InnerNode {
+        let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
+        InnerNode { left: child, right: child }
+    }
+
     fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
         if value == Self::EMPTY_VALUE {
             self.leaves.remove(&key.value())
@@ -319,6 +324,11 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         }
     }
 
+    fn maybe_get_leaf(&self, key: &LeafIndex<DEPTH>) -> Option<Word> {
+        let leaf_pos = key.value();
+        self.leaves.get(&leaf_pos).copied()
+    }
+
     fn hash_leaf(leaf: &Word) -> RpoDigest {
         // `SimpleSmt` takes the leaf value itself as the hash
         leaf.into()
@@ -328,6 +338,15 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
         Self::hash_leaf(value)
     }
 
+    fn get_prospective_leaf(
+        &self,
+        _existing_leaf: Option<&Word>,
+        _key: &LeafIndex<DEPTH>,
+        value: &Word,
+    ) -> Word {
+        *value
+    }
+
     fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
         *key
     }