diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 104ceb4..688ef16 100644 --- a/src/merkle/index.rs +++ b/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. diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index b1f0839..866aa19 100644 --- a/src/merkle/smt/full/mod.rs +++ b/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 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::::from(*key).value(); @@ -281,6 +288,11 @@ impl SparseMerkleTree for Smt { } } + fn maybe_get_leaf(&self, key: &RpoDigest) -> Option { + let leaf_pos = LeafIndex::::from(*key).value(); + self.leaves.get(&leaf_pos).cloned() + } + fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest { leaf.hash() } @@ -304,6 +316,27 @@ impl SparseMerkleTree for Smt { } } + fn get_prospective_leaf( + &self, + existing_leaf: Option<&SmtLeaf>, + key: &RpoDigest, + value: &Word, + ) -> SmtLeaf { + //let leaf_index: LeafIndex = 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 { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index dc3faf3..b692cf6 100644 --- a/src/merkle/smt/full/tests.rs +++ b/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 diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 790e317..11b1ba7 100644 --- a/src/merkle/smt/mod.rs +++ b/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) -> &'static str { + std::any::type_name::() + } + 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 { +pub(crate) trait SparseMerkleTree: 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 { 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 { return Ok(value); } - // Compute the nodes we'll need to make and remove. - let mut removals: Vec = Vec::with_capacity(DEPTH as usize); - let mut additions: Vec<(NodeIndex, InnerNode)> = Vec::with_capacity(DEPTH as usize); + //let mut mutations = Mutations::new(); + //self.prospective_insert(&mut mutations, key.clone(), value.clone()); + //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, + // 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 = 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, + 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 = 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 { } }; + 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 - // default value. - removals.push(node_index); + mutations.maybe_remove_node(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 { - return Err(MerkleError::ConflictingRoots(vec![expected_root, new_root])); + assert_eq!(control.root(), mutations.new_root); + } + + fn start_muts(&self) -> Muts { + Muts:: { + //mutations: Default::default(), + node_mutations: Default::default(), + pair_insertions: Default::default(), + new_root: self.root(), + } + } + + //fn apply_mutations(&mut self, mut mutations: Mutations) { + // 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) { + 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); + } } - // Actual mutations start here. - - self.insert_value(key, value); - - for index in removals.drain(..) { - self.remove_inner_node(index); + 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 (index, new_node) in additions.drain(..) { - self.insert_inner_node(index, new_node); + for (key, value) in pair_insertions.into_iter() { + self.insert_value(key, value); } - 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 { /// 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; + 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 { /// 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; + /// Returns the hash of a leaf fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; @@ -264,6 +603,13 @@ pub(crate) trait SparseMerkleTree { /// 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; @@ -341,3 +687,340 @@ impl TryFrom for LeafIndex { Self::new(node_index.value()) } } + +// MUTATIONS +//#[derive(Debug, Clone, PartialEq)] +//#[cfg(not(test))] +//pub enum Mutation { +// 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 { + //mutations: Vec>, + node_mutations: BTreeMap, + pair_insertions: Vec<(K, V)>, + new_root: RpoDigest, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct MaybeNode { + inner: Option, + index: NodeIndex, +} +impl MaybeNode { + 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 Muts { + fn maybe_get_values(&self, key: &K) -> Vec { + 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(&self, tree: &T, key: &K) -> Option + where + T: SparseMerkleTree, + { + 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 = 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 { + 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 { + use NodeMutation::*; + match self.node_mutations.get(&index) { + Some(Addition(node)) => MaybeNode::from(node.clone(), index), + Some(Removal) => { + MaybeNode::from(MaybeNode:: { 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 { +// node_removals: Vec, +// node_additions: Vec<(NodeIndex, InnerNode)>, +// pair_insertions: Vec<(K, V)>, +// new_root: RpoDigest, +//} +// +//impl Mutations { +// 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 { +// 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) +// } +//} diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 8e41b55..4e1c7e6 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -299,6 +299,11 @@ impl SparseMerkleTree for SimpleSmt { 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, value: Word) -> Option { if value == Self::EMPTY_VALUE { self.leaves.remove(&key.value()) @@ -319,6 +324,11 @@ impl SparseMerkleTree for SimpleSmt { } } + fn maybe_get_leaf(&self, key: &LeafIndex) -> Option { + 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 SparseMerkleTree for SimpleSmt { Self::hash_leaf(value) } + fn get_prospective_leaf( + &self, + _existing_leaf: Option<&Word>, + _key: &LeafIndex, + value: &Word, + ) -> Word { + *value + } + fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key }