diff --git a/Cargo.lock b/Cargo.lock index b88f06c03..7eb6f2b6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3976,6 +3976,7 @@ dependencies = [ "quickcheck_macros", "rand 0.8.6", "serde", + "serde_json", ] [[package]] @@ -4561,8 +4562,8 @@ dependencies = [ "logos-blockchain-poseidon2", "num-bigint", "rand 0.8.6", - "rpds", "serde", + "serde_json", ] [[package]] diff --git a/merkle/blake2btree/Cargo.toml b/merkle/blake2btree/Cargo.toml index f1308701f..c13a8aa39 100644 --- a/merkle/blake2btree/Cargo.toml +++ b/merkle/blake2btree/Cargo.toml @@ -19,6 +19,7 @@ quickcheck = { workspace = true } quickcheck_macros = { workspace = true } rand = { features = ["std", "std_rng"], workspace = true } serde = { features = ["derive"], workspace = true } +serde_json = { features = ["alloc"], workspace = true } [lints] workspace = true diff --git a/merkle/blake2btree/src/lib.rs b/merkle/blake2btree/src/lib.rs index 38c222c35..ebee10a9d 100644 --- a/merkle/blake2btree/src/lib.rs +++ b/merkle/blake2btree/src/lib.rs @@ -421,6 +421,115 @@ mod tests { assert!(tree.path(&TestLeaf::from_usize(2)).is_none()); } + #[test] + fn test_sparse_recovery_preserves_paths_and_supports_mutation() { + let capacity = 1usize << lb_dynamic_merkle::TREE_HEIGHT_EXCEPT_ROOT; + let entries = [ + (1, TestLeaf::from_usize(1), TestLeaf::from_usize(101)), + (7, TestLeaf::from_usize(7), TestLeaf::from_usize(107)), + ( + 1usize << 31, + TestLeaf::from_usize(31), + TestLeaf::from_usize(131), + ), + ( + capacity - 1, + TestLeaf::from_usize(32), + TestLeaf::from_usize(132), + ), + ]; + let serialized = serialized_tree(&entries); + let tree: Blake2bTree = + serde_json::from_str(&serialized).expect("sparse tree should deserialize"); + + assert_eq!(tree.size(), entries.len()); + for (_, key, item) in &entries { + verify_path(&tree, key, item); + } + + let inserted_key = TestLeaf::from_usize(99); + let inserted_item = TestLeaf::from_usize(199); + let (tree, position) = tree.insert(inserted_key, inserted_item); + assert_eq!(position, 0); + for (_, key, item) in &entries { + verify_path(&tree, key, item); + } + verify_path(&tree, &inserted_key, &inserted_item); + + let removed_key = entries[3].1; + let (tree, removed_item) = tree.remove(&removed_key).unwrap(); + assert_eq!(removed_item, entries[3].2); + assert!(tree.path(&removed_key).is_none()); + + let (tree, position) = tree.insert(removed_key, removed_item); + assert_eq!(position, 2); + for (_, key, item) in &entries { + verify_path(&tree, key, item); + } + verify_path(&tree, &inserted_key, &inserted_item); + } + + #[test] + fn test_public_deserialization_rejects_out_of_capacity_position() { + let capacity = 1usize << lb_dynamic_merkle::TREE_HEIGHT_EXCEPT_ROOT; + let serialized = + serialized_tree(&[(capacity, TestLeaf::from_usize(1), TestLeaf::from_usize(1))]); + + let error = serde_json::from_str::>(&serialized) + .expect_err("out-of-capacity positions must be rejected"); + assert!(error.to_string().contains("exceeds capacity")); + } + + #[test] + fn test_public_deserialization_rejects_duplicate_logical_keys() { + let key = TestLeaf::from_usize(1); + let serialized = serialized_tree(&[ + (0, key, TestLeaf::from_usize(101)), + (1, key, TestLeaf::from_usize(102)), + ]); + + let error = serde_json::from_str::>(&serialized) + .expect_err("duplicate logical keys must be rejected"); + assert!(error.to_string().contains("duplicate key at position 1")); + } + + #[test] + fn test_public_deserialization_rejects_duplicate_serialized_positions() { + let value = + serde_json::to_string(&(TestLeaf::from_usize(1), TestLeaf::from_usize(101))).unwrap(); + let serialized = format!("{{\"0\":{value},\"0\":{value}}}"); + + let error = serde_json::from_str::>(&serialized) + .expect_err("duplicate serialized positions must be rejected"); + assert!(error.to_string().contains("duplicate positions")); + } + + fn serialized_tree(entries: &[(usize, TestLeaf, TestLeaf)]) -> String { + let entries = entries + .iter() + .map(|(position, key, item)| { + format!( + "{}:{}", + serde_json::to_string(&position.to_string()).unwrap(), + serde_json::to_string(&(key, item)).unwrap() + ) + }) + .collect::>() + .join(","); + format!("{{{entries}}}") + } + + fn verify_path(tree: &Blake2bTree, key: &TestLeaf, item: &TestLeaf) { + let mut current = Blake2bLeaf::leaf(key, item); + for node in tree.path(key).expect("path should exist") { + current = match node { + MerkleNode::Left(sibling) => Blake2bMerkleHasher::compress(&sibling, ¤t), + MerkleNode::Right(sibling) => Blake2bMerkleHasher::compress(¤t, &sibling), + }; + } + assert_eq!(current, tree.root()); + } + // `Blake2bTree` is a type alias for the foreign `MerkleTree`, so `Arbitrary` // (also foreign) can't be implemented on it directly. Wrap it in a local // newtype for the property test. @@ -457,7 +566,7 @@ mod tests { let compressed = original_tree.compressed(); // Recover the tree from compressed format - let recovered_tree: Blake2bTree<_, _> = compressed.into(); + let recovered_tree: Blake2bTree<_, _> = compressed.try_into().unwrap(); recovered_tree == original_tree && recovered_tree.root() == original_tree.root() } diff --git a/merkle/dynamic-merkle/Cargo.toml b/merkle/dynamic-merkle/Cargo.toml index ced714fe5..4479a4b18 100644 --- a/merkle/dynamic-merkle/Cargo.toml +++ b/merkle/dynamic-merkle/Cargo.toml @@ -10,7 +10,6 @@ repository.workspace = true version.workspace = true [dependencies] -rpds = { features = ["serde"], workspace = true } serde = { features = ["alloc", "derive", "rc"], workspace = true } [dev-dependencies] @@ -18,6 +17,7 @@ ark-ff = { workspace = true } lb-poseidon2 = { workspace = true } num-bigint = { workspace = true } rand = { workspace = true } +serde_json = { features = ["alloc"], workspace = true } [lints] workspace = true diff --git a/merkle/dynamic-merkle/src/lib.rs b/merkle/dynamic-merkle/src/lib.rs index da8476277..dcb1844a4 100644 --- a/merkle/dynamic-merkle/src/lib.rs +++ b/merkle/dynamic-merkle/src/lib.rs @@ -15,8 +15,6 @@ use std::{fmt, marker::PhantomData, sync::Arc}; -use rpds::RedBlackTreeSetSync; - /// Abstraction over the hash type and hashing operations a /// [`DynamicMerkleTree`] needs. /// @@ -93,35 +91,25 @@ enum Node { left_subtree_size: usize, height: usize, }, - // An empty inner node, representing an unexpanded empty subtree, to avoid - // allocating a full subtree when not necessary. - // Can only be found in the right subtree of an inner node. + // An unexpanded, fully-empty subtree. Height zero represents one empty + // leaf position; larger heights compactly represent multiple empty + // leaves, avoiding allocations for ranges that are not occupied. Empty { height: usize, }, - // A leaf node (possibly) holding a hash, will be empty after a removal + // An occupied leaf node. Leaf { - value: Option, + value: Hash, }, } fn hash(left: &Node, right: &Node) -> H::Hash { - let left = match left { - Node::Inner { value, .. } => *value, - Node::Leaf { value } => (*value).unwrap_or(H::EMPTY_VALUE), - Node::Empty { .. } => panic!("Empty node in left subtree is not allowed"), - }; - let right = match right { - Node::Inner { value, .. } => *value, - Node::Leaf { value } => (*value).unwrap_or(H::EMPTY_VALUE), - Node::Empty { height } => H::empty_subtree_root(*height), - }; - H::compress(&left, &right) + H::compress(&left.value::(), &right.value::()) } impl Node { const fn new(value: Hash) -> Self { - Self::Leaf { value: Some(value) } + Self::Leaf { value } } fn size(&self) -> usize { @@ -131,8 +119,8 @@ impl Node { right_subtree_size, .. } => left_subtree_size + right_subtree_size, - Self::Leaf { value: Some(_) } => 1, - Self::Empty { .. } | Self::Leaf { value: None } => 0, + Self::Leaf { .. } => 1, + Self::Empty { .. } => 0, } } @@ -141,6 +129,24 @@ impl Node { 1 << self.height() } + fn first_empty_index(&self) -> Option { + match self { + Self::Inner { left, right, .. } => { + if left.size() < left.capacity() { + left.first_empty_index() + } else if right.size() < right.capacity() { + right + .first_empty_index() + .map(|index| left.capacity() + index) + } else { + None + } + } + Self::Empty { .. } => Some(0), + Self::Leaf { .. } => None, + } + } + const fn height(&self) -> usize { match self { Self::Inner { height, .. } | Self::Empty { height } => *height, @@ -150,14 +156,73 @@ impl Node { } impl Node { + /// Recursively validates a deserialized node tree and rebuilds derived + /// inner-node state. + /// + /// Inner-node hashes, subtree sizes, and heights are recomputed from the + /// children. Fully empty sibling subtrees are collapsed into a single + /// empty parent. Invalid subtree heights or mismatched sibling heights + /// are rejected. + fn rebuild_and_validate(node: &Arc) -> Result, &'static str> + where + H: MerkleHasher, + { + let node = match node.as_ref() { + Self::Empty { height } => { + if *height > TREE_HEIGHT_EXCEPT_ROOT { + return Err("empty subtree height exceeds tree height"); + } + return Ok(Arc::clone(node)); + } + Self::Leaf { .. } => return Ok(Arc::clone(node)), + Self::Inner { left, right, .. } => { + let left = Self::rebuild_and_validate::(left)?; + let right = Self::rebuild_and_validate::(right)?; + // Malformed serialized input must be rejected as a serde + // error; `new_inner`'s assertion is for internal misuse. + if left.height() != right.height() { + return Err("sibling subtrees have different heights"); + } + Self::new_inner::(left, right) + } + }; + if node.height() > TREE_HEIGHT_EXCEPT_ROOT { + return Err("tree height exceeds fixed tree height"); + } + Ok(Arc::new(node)) + } + fn new_inner(left: Arc, right: Arc) -> Self where H: MerkleHasher, { + if let ( + Self::Empty { + height: left_height, + }, + Self::Empty { + height: right_height, + }, + ) = (left.as_ref(), right.as_ref()) + { + assert_eq!( + left_height, right_height, + "empty sibling subtrees must have equal heights" + ); + return Self::Empty { + height: left_height + 1, + }; + } + + assert_eq!( + left.height(), + right.height(), + "sibling subtrees must have equal heights" + ); Self::Inner { right_subtree_size: right.size(), left_subtree_size: left.size(), - height: left.height().max(right.height()) + 1, + height: left.height() + 1, value: hash::(&left, &right), left, right, @@ -221,9 +286,9 @@ impl Node { H: MerkleHasher, { self.insert_or_modify::(index, |node| match node { - Self::Leaf { value: None } | Self::Empty { .. } => Self::new(value), - Self::Leaf { value: Some(_) } => panic!("Cannot insert into a non-empty leaf node"), - _ => panic!("Cannot insert into a non-terminal node"), + Self::Empty { .. } => Self::new(value), + Self::Leaf { .. } => panic!("Cannot insert into a non-empty leaf node"), + Self::Inner { .. } => panic!("Cannot insert into a non-terminal node"), }) } @@ -232,7 +297,7 @@ impl Node { H: MerkleHasher, { self.insert_or_modify::(index, move |node| match node { - Self::Leaf { value: Some(_) } => Self::Leaf { value: None }, + Self::Leaf { .. } => Self::Empty { height: 0 }, _ => panic!("Cannot remove from a empty / non-leaf node"), }) } @@ -242,7 +307,7 @@ impl Node { H: MerkleHasher, { self.insert_or_modify::(index, |node| match node { - Self::Leaf { value: Some(_) } => Self::new(value), + Self::Leaf { .. } => Self::new(value), _ => panic!("Cannot update an empty / non-leaf node"), }) } @@ -281,8 +346,8 @@ impl Node { Some(path) } } - Self::Leaf { value: Some(_) } => Some(Vec::new()), - Self::Leaf { value: None } | Self::Empty { .. } => None, + Self::Leaf { .. } => Some(Vec::new()), + Self::Empty { .. } => None, } } @@ -291,8 +356,7 @@ impl Node { H: MerkleHasher, { match self { - Self::Inner { value, .. } | Self::Leaf { value: Some(value) } => *value, - Self::Leaf { value: None } => H::EMPTY_VALUE, + Self::Inner { value, .. } | Self::Leaf { value } => *value, Self::Empty { height } => H::empty_subtree_root(*height), } } @@ -301,14 +365,14 @@ impl Node { /// A dynamic persistent Merkle tree that supports insertion and removal of /// leaf hashes. /// -/// Removed leaves are replaced with an empty leaf node, which prevents -/// the whole tree reordering and their position is recorded for future -/// insertions. Compared to a MPT, the height of this tree is predictable and -/// bounded by the number of items, for example allowing for efficient and -/// simple proof of memberships for `PoL`. +/// `Node::Empty` is the sole representation of free space: height zero +/// represents one free leaf and larger heights compactly represent free +/// ranges. Sparse recovery and normal mutation therefore produce the same +/// canonical structure. Compared to a MPT, the height of this tree is +/// predictable and bounded by the number of items, for example allowing for +/// efficient and simple proof of memberships for `PoL`. pub struct DynamicMerkleTree { root: Arc>, - holes: RedBlackTreeSetSync, _hasher: PhantomData, } @@ -316,7 +380,6 @@ impl Clone for DynamicMerkleTree { fn clone(&self) -> Self { Self { root: Arc::clone(&self.root), - holes: self.holes.clone(), _hasher: PhantomData, } } @@ -329,19 +392,16 @@ where fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("DynamicMerkleTree") .field("root", &self.root) - .field("holes", &self.holes) .finish() } } impl Default for DynamicMerkleTree { fn default() -> Self { - let holes = RedBlackTreeSetSync::new_sync(); Self { root: Arc::new(Node::Empty { height: TREE_HEIGHT_EXCEPT_ROOT, }), - holes, _hasher: PhantomData, } } @@ -364,8 +424,9 @@ impl DynamicMerkleTree { /// Inserts the leaf hash `value` at the lowest available position and /// returns the updated tree together with the index it was assigned. /// - /// Positions freed by [`remove`](Self::remove) are reused before the tree - /// grows, so the smallest free index is always chosen. + /// The lowest available position is derived solely from the tree + /// structure, so sparse recovery and removals use the same representation + /// and the smallest free index is always chosen. /// /// The original tree is left unchanged (the structure is persistent). /// @@ -379,16 +440,14 @@ impl DynamicMerkleTree { "max capacity reached, cannot insert more items" ); - let (holes, index) = self.holes.first().map_or_else( - || (self.holes.clone(), self.root.size()), - |hole| (self.holes.remove(hole), *hole), - ); - + let index = self + .root + .first_empty_index() + .expect("tree has capacity but no empty position"); let root = self.root.insert_at::(index, value); ( Self { root, - holes, _hasher: PhantomData, }, index, @@ -397,9 +456,9 @@ impl DynamicMerkleTree { /// Removes the leaf at `index`, returning the updated tree. /// - /// The leaf is replaced with an empty one and its position is recorded as a - /// hole for reuse by a future [`insert`](Self::insert); the tree is not - /// otherwise restructured. The original tree is left unchanged. + /// The leaf is replaced with `Node::Empty { height: 0 }`. Empty sibling + /// subtrees are collapsed while the ancestors are rebuilt, so the tree + /// remains canonical. The original tree is left unchanged. /// /// # Panics /// @@ -410,10 +469,8 @@ impl DynamicMerkleTree { assert!(index < self.root.capacity(), "Index out of bounds"); let root = self.root.remove_at::(index); - let holes = self.holes.insert(index); Self { root, - holes, _hasher: PhantomData, } } @@ -421,8 +478,8 @@ impl DynamicMerkleTree { /// Replaces the leaf hash at `index`, returning the updated tree. /// /// Unlike [`remove`](Self::remove), the leaf is replaced with another value - /// instead of being emptied, so the position is neither freed nor recorded - /// as a hole. The original tree is left unchanged. + /// instead of being emptied, so the position remains occupied. The + /// original tree is left unchanged. /// /// # Panics /// @@ -435,7 +492,6 @@ impl DynamicMerkleTree { let root = self.root.update_at::(index, value); Self { root, - holes: self.holes.clone(), _hasher: PhantomData, } } @@ -462,8 +518,10 @@ impl DynamicMerkleTree { self.root.path::(index)?.try_into().ok() } - /// Rebuilds a tree placing each leaf hash at its given index, filling the - /// gaps between indices with holes. + /// Rebuilds a tree placing each leaf hash at its given index, representing + /// gaps as empty subtrees rather than materializing individual empty + /// leaves. Empty sibling subtrees are represented by their single empty + /// parent, matching the canonical structure produced by mutation. /// /// The values must be yielded in strictly increasing index order; this is /// the inverse of enumerating a tree's occupied positions and is meant @@ -475,43 +533,72 @@ impl DynamicMerkleTree { /// bounds. #[must_use] pub fn from_sorted_items(items: impl IntoIterator) -> Self { - let mut tree = Self::new(); - let mut current_pos = 0; - for (pos, value) in items { - while current_pos < pos { - // Insert a hole for the missing position - tree = tree.insert_hole(current_pos); - current_pos += 1; - } - - tree.root = tree.root.insert_at::(pos, value); - current_pos = pos + 1; - } - tree - } - - // This is only for maintaining holes information when recovering - // the tree from a compressed format, should not be used otherwise. - fn insert_hole(&self, index: usize) -> Self { + let mut items = items.into_iter().peekable(); + let root = Self::build_sparse_subtree(&mut items, 0, TREE_HEIGHT_EXCEPT_ROOT); assert!( - index < self.root.capacity(), - "Index out of bounds for inserting an empty node" + items.next().is_none(), + "all indices must be within tree capacity" ); - - let holes = self.holes.insert(index); - let root = self - .root - .insert_or_modify::(index, |node| match node { - Node::Empty { .. } => Node::Leaf { value: None }, - _ => panic!("Cannot insert a hole into a non-empty/non-leaf node"), - }); - Self { root, - holes, _hasher: PhantomData, } } + + /// Builds a sparse subtree starting at `subtree_start`, representing + /// unoccupied ranges as [`Node::Empty`] rather than materializing + /// individual empty leaves. + /// + /// A subtree of `height` covers the half-open range + /// `[subtree_start, subtree_start + 2^height)`. Thus, `subtree_capacity = + /// 2^height` is the width of this subtree, not an absolute index bound, + /// and `midpoint = subtree_start + subtree_capacity / 2` divides the range + /// into equal left and right halves. + /// + /// The iterator must yield entries in strictly increasing absolute + /// position order. At `height == 0`, this subtree represents exactly one + /// leaf position, so any consumed item is at `subtree_start`. The function + /// consumes only entries belonging to the current subtree. + fn build_sparse_subtree( + items: &mut std::iter::Peekable, + subtree_start: usize, + height: usize, + ) -> Arc> + where + I: Iterator, + { + let Some(&(position, _)) = items.peek() else { + return Arc::new(Node::Empty { height }); + }; + let subtree_capacity = 1usize << height; + assert!( + position >= subtree_start && position - subtree_start < subtree_capacity, + "indices must be strictly increasing and within bounds" + ); + if height == 0 { + let (_, value) = items.next().expect("peeked item must be available"); + return Arc::new(Node::new(value)); + } + + let midpoint = subtree_start + (subtree_capacity / 2); + let left = if items + .peek() + .is_some_and(|(position, _)| *position < midpoint) + { + Self::build_sparse_subtree(items, subtree_start, height - 1) + } else { + Arc::new(Node::Empty { height: height - 1 }) + }; + let right = if items + .peek() + .is_some_and(|(position, _)| *position < subtree_start + subtree_capacity) + { + Self::build_sparse_subtree(items, midpoint, height - 1) + } else { + Arc::new(Node::Empty { height: height - 1 }) + }; + Arc::new(Node::new_inner::(left, right)) + } } impl PartialEq for DynamicMerkleTree { @@ -524,22 +611,22 @@ impl Eq for DynamicMerkleTree {} /// [`serde`](::serde) support for [`DynamicMerkleTree`]. /// -/// The tree serializes as its root node and the set of holes; on -/// deserialization the two are reassembled into a tree. Requires the hasher's +/// The tree serializes its root node. Deserialization recursively validates +/// and rebuilds the node tree, recomputing derived inner-node state and +/// collapsing fully empty sibling subtrees. Requires the hasher's /// [`Hash`](MerkleHasher::Hash) type to implement the corresponding `serde` /// traits. pub mod serde { - use std::{marker::PhantomData, sync::Arc}; + use std::sync::Arc; - use rpds::RedBlackTreeSetSync; use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeStruct as _}; use super::MerkleHasher; #[derive(Deserialize)] + #[serde(deny_unknown_fields)] struct Raw { root: Arc>, - holes: RedBlackTreeSetSync, } impl Serialize for super::DynamicMerkleTree @@ -551,9 +638,8 @@ pub mod serde { where S: Serializer, { - let mut state = serializer.serialize_struct("DynamicMerkleTree", 2)?; + let mut state = serializer.serialize_struct("DynamicMerkleTree", 1)?; state.serialize_field("root", &self.root)?; - state.serialize_field("holes", &self.holes)?; state.end() } } @@ -568,10 +654,18 @@ pub mod serde { D: Deserializer<'de>, { let raw = Raw::::deserialize(deserializer)?; + let root = super::Node::rebuild_and_validate::(&raw.root) + .map_err(serde::de::Error::custom)?; + if root.height() != super::TREE_HEIGHT_EXCEPT_ROOT + || matches!(root.as_ref(), super::Node::Leaf { .. }) + { + return Err(serde::de::Error::custom( + "dynamic Merkle tree root has an invalid shape", + )); + } Ok(Self { - root: raw.root, - holes: raw.holes, - _hasher: PhantomData, + root, + _hasher: std::marker::PhantomData, }) } } @@ -653,7 +747,7 @@ mod tests { } #[test] - fn test_hole_management() { + fn test_free_position_management() { let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let mut rng = rand::thread_rng(); let a = fr_from_rng(&mut rng); @@ -745,6 +839,169 @@ mod tests { assert_eq!(tree.size(), 3); } + #[test] + fn test_sparse_recovery_does_not_materialize_gaps() { + let value = fr_from_usize(1); + let recovered = + DynamicMerkleTree::::from_sorted_items([((1usize << 31) + 1, value)]); + + assert_eq!(recovered.size(), 1); + let (recovered, index) = recovered.insert(fr_from_usize(2)); + assert_eq!(index, 0); + assert_eq!(recovered.size(), 2); + } + + #[test] + fn test_single_removal_uses_empty_leaf() { + let tree = DynamicMerkleTree::::new(); + let (tree, first_position) = tree.insert(fr_from_usize(1)); + let (tree, second_position) = tree.insert(fr_from_usize(2)); + + let removed = tree.remove(first_position); + + assert_eq!(removed.size(), 1); + assert!(matches!( + subtree_at(&removed.root, first_position, 0), + Node::Empty { height: 0 } + )); + assert!(removed.path(first_position).is_none()); + assert!(removed.path(second_position).is_some()); + assert_canonical(&removed.root); + } + + #[test] + fn test_sibling_empty_subtrees_collapse() { + let mut tree = DynamicMerkleTree::::new(); + for value in 0..3 { + tree = tree.insert(fr_from_usize(value)).0; + } + + tree = tree.remove(0); + tree = tree.remove(1); + + assert!(matches!( + subtree_at(&tree.root, 0, 1), + Node::Empty { height: 1 } + )); + assert_eq!(tree.size(), 1); + assert_canonical(&tree.root); + } + + #[test] + fn test_empty_subtrees_collapse_recursively() { + let mut tree = DynamicMerkleTree::::new(); + for value in 0..5 { + tree = tree.insert(fr_from_usize(value)).0; + } + + for position in 0..4 { + tree = tree.remove(position); + } + + assert!(matches!( + subtree_at(&tree.root, 0, 2), + Node::Empty { height: 2 } + )); + assert_eq!(tree.size(), 1); + assert_canonical(&tree.root); + } + + #[test] + fn test_removing_every_item_restores_canonical_empty_tree() { + let mut tree = DynamicMerkleTree::::new(); + for value in 0..8 { + tree = tree.insert(fr_from_usize(value)).0; + } + for position in 0..8 { + tree = tree.remove(position); + } + + assert!(matches!( + tree.root.as_ref(), + Node::Empty { + height: TREE_HEIGHT_EXCEPT_ROOT + } + )); + assert_eq!(tree.root(), DynamicMerkleTree::::new().root()); + assert_canonical(&tree.root); + } + + #[test] + fn test_remove_and_reinsert_restores_root() { + let tree = DynamicMerkleTree::::new(); + let value = fr_from_usize(1); + let (tree, _) = tree.insert(value); + let (tree, _) = tree.insert(fr_from_usize(2)); + let original_root = tree.root(); + + let tree = tree.remove(0); + let (tree, position) = tree.insert(value); + + assert_eq!(position, 0); + assert_eq!(tree.root(), original_root); + assert_canonical(&tree.root); + } + + #[test] + fn test_mutation_and_sparse_recovery_have_same_structure() { + let capacity = 1usize << TREE_HEIGHT_EXCEPT_ROOT; + let items = [ + (1, fr_from_usize(1)), + (7, fr_from_usize(7)), + (1usize << 31, fr_from_usize(31)), + (capacity - 1, fr_from_usize(32)), + ]; + let recovered = DynamicMerkleTree::::from_sorted_items(items); + let mutated = tree_from_indexed_mutation(items); + + assert_eq!(mutated.size(), items.len()); + assert_eq!(mutated.root, recovered.root); + assert_canonical(&mutated.root); + assert_canonical(&recovered.root); + } + + #[test] + fn test_direct_serde_roundtrip_preserves_canonical_structure() { + let items = [ + (1, [1; 32]), + (7, [7; 32]), + (1usize << 31, [31; 32]), + ((1usize << 32) - 1, [32; 32]), + ]; + let tree = DynamicMerkleTree::::from_sorted_items(items); + let (tree, inserted_position) = tree.insert([99; 32]); + let tree = tree.remove(inserted_position); + let serialized = serde_json::to_string(&tree).expect("serialize dynamic tree"); + let recovered: DynamicMerkleTree = + serde_json::from_str(&serialized).expect("deserialize dynamic tree"); + + assert!(!serialized.contains("holes")); + assert_eq!(tree.root, recovered.root); + assert_canonical(&recovered.root); + } + + #[test] + fn test_direct_serde_canonicalizes_empty_inner_nodes() { + #[derive(::serde::Serialize)] + struct RawTree { + root: Arc>, + } + + let raw = RawTree { + root: noncanonical_empty_tree::(TREE_HEIGHT_EXCEPT_ROOT), + }; + let serialized = serde_json::to_string(&raw).expect("serialize non-canonical tree"); + let recovered: DynamicMerkleTree = + serde_json::from_str(&serialized).expect("deserialize non-canonical tree"); + + assert!(matches!( + recovered.root.as_ref(), + Node::Empty { + height: TREE_HEIGHT_EXCEPT_ROOT + } + )); + } + #[test] fn test_remove_single_item() { let tree: DynamicMerkleTree = DynamicMerkleTree::new(); @@ -795,7 +1052,7 @@ mod tests { } #[test] - fn test_smallest_hole_selection() { + fn test_smallest_free_position_selection() { let tree: DynamicMerkleTree = DynamicMerkleTree::new(); // Insert items at positions 0, 1, 2, 3, 4 @@ -805,23 +1062,23 @@ mod tests { let (tree, _) = tree.insert(fr_from_rng(&mut rand::thread_rng())); let (tree, _) = tree.insert(fr_from_rng(&mut rand::thread_rng())); - // Remove items at positions 3, 1, 4 (creating holes in that order) + // Remove items at positions 3, 1, 4 (creating free positions in that order) let tree = tree.remove(3); let tree = tree.remove(1); let tree = tree.remove(4); - // Now we have holes at positions 1, 3, 4 - // The smallest hole should be selected first (position 1) + // Now we have free positions at 1, 3, 4. The smallest free position + // should be selected first (position 1). let (tree, index1) = tree.insert(fr_from_rng(&mut rand::thread_rng())); - assert_eq!(index1, 1, "Should select smallest hole first"); + assert_eq!(index1, 1, "Should select smallest free position first"); - // Next insertion should use the next smallest hole (position 3) + // Next insertion should use the next smallest free position (position 3) let (tree, index2) = tree.insert(fr_from_rng(&mut rand::thread_rng())); - assert_eq!(index2, 3, "Should select next smallest hole"); + assert_eq!(index2, 3, "Should select next smallest free position"); - // Final insertion should use the last hole (position 4) + // Final insertion should use the last free position (position 4) let (_, index3) = tree.insert(fr_from_rng(&mut rand::thread_rng())); - assert_eq!(index3, 4, "Should select remaining hole"); + assert_eq!(index3, 4, "Should select remaining free position"); } #[test] @@ -911,4 +1168,137 @@ mod tests { "Computed root from path doesn't match expected root" ); } + + fn assert_canonical(node: &Arc>) { + match node.as_ref() { + Node::Empty { .. } | Node::Leaf { .. } => {} + Node::Inner { + left, + right, + left_subtree_size, + right_subtree_size, + height, + .. + } => { + assert_eq!(left.height(), right.height()); + assert!(!matches!( + (left.as_ref(), right.as_ref()), + (Node::Empty { .. }, Node::Empty { .. }) + )); + assert_eq!(*height, left.height() + 1); + assert_eq!(*left_subtree_size, left.size()); + assert_eq!(*right_subtree_size, right.size()); + assert_canonical(left); + assert_canonical(right); + } + } + } + + fn subtree_at(node: &Arc>, index: usize, target_height: usize) -> &Node { + assert!(target_height <= node.height()); + if target_height == node.height() { + return node; + } + match node.as_ref() { + Node::Inner { left, right, .. } => { + if index < left.capacity() { + subtree_at(left, index, target_height) + } else { + subtree_at(right, index - left.capacity(), target_height) + } + } + Node::Empty { .. } => panic!("empty subtree is already canonical at this range"), + Node::Leaf { .. } => panic!("leaf is already at the smallest range"), + } + } + + fn tree_from_indexed_mutation( + items: impl IntoIterator, + ) -> DynamicMerkleTree { + let root = items.into_iter().fold( + Arc::new(Node::Empty { + height: TREE_HEIGHT_EXCEPT_ROOT, + }), + |root, (position, value)| insert_at_any_position(&root, position, value), + ); + DynamicMerkleTree { + root, + _hasher: PhantomData, + } + } + + fn insert_at_any_position(node: &Arc>, index: usize, value: Fr) -> Arc> { + match node.as_ref() { + Node::Inner { left, right, .. } => { + if index < left.capacity() { + Arc::new(Node::new_inner::( + insert_at_any_position(left, index, value), + Arc::clone(right), + )) + } else { + Arc::new(Node::new_inner::( + Arc::clone(left), + insert_at_any_position(right, index - left.capacity(), value), + )) + } + } + Node::Empty { height } if *height > 0 => { + let half = 1usize << (height - 1); + let left = Arc::new(Node::Empty { height: height - 1 }); + let right = Arc::new(Node::Empty { height: height - 1 }); + if index < half { + Arc::new(Node::new_inner::( + insert_at_any_position(&left, index, value), + right, + )) + } else { + Arc::new(Node::new_inner::( + left, + insert_at_any_position(&right, index - half, value), + )) + } + } + Node::Empty { .. } => Arc::new(Node::new(value)), + Node::Leaf { .. } => panic!("cannot insert into an occupied position"), + } + } + + fn noncanonical_empty_tree(height: usize) -> Arc> + where + H: MerkleHasher, + { + assert!(height > 0); + let left = if height == 1 { + Arc::new(Node::Empty { height: 0 }) + } else { + noncanonical_empty_tree::(height - 1) + }; + let right = Arc::new(Node::Empty { height: height - 1 }); + Arc::new(Node::Inner { + value: hash::(&left, &right), + left_subtree_size: left.size(), + right_subtree_size: right.size(), + height, + left, + right, + }) + } + + struct SerdeHasher; + + impl MerkleHasher for SerdeHasher { + type Hash = [u8; 32]; + + const EMPTY_VALUE: Self::Hash = [0; 32]; + + fn compress(left: &Self::Hash, right: &Self::Hash) -> Self::Hash { + let mut hash = [0; 32]; + for (output, (left, right)) in hash.iter_mut().zip(left.iter().zip(right)) { + *output = left.wrapping_add(*right); + } + hash + } + + empty_subtree_root!([u8; 32]); + } } diff --git a/merkle/tree/src/lib.rs b/merkle/tree/src/lib.rs index 719d93712..6558a559b 100644 --- a/merkle/tree/src/lib.rs +++ b/merkle/tree/src/lib.rs @@ -1,4 +1,8 @@ -use std::{collections::BTreeMap, fmt, marker::PhantomData}; +use std::{ + collections::{BTreeMap, HashSet}, + fmt, + marker::PhantomData, +}; use lb_dynamic_merkle::MerkleHasher; pub use lb_dynamic_merkle::{DynamicMerkleTree, MerkleNode, MerklePath}; @@ -268,13 +272,39 @@ where } } -impl From> for MerkleTree +#[derive(Error, Debug, Clone, PartialEq, Eq)] +pub enum RecoveryError { + #[error("compressed Merkle tree position {position} exceeds capacity {capacity}")] + PositionOutOfBounds { position: usize, capacity: usize }, + #[error("compressed Merkle tree contains a duplicate key at position {position}")] + DuplicateKey { position: usize }, +} + +impl TryFrom> for MerkleTree where Key: Clone + std::hash::Hash + Eq, Item: Clone, Leaf: LeafExtractor, { - fn from(compressed: CompressedMerkleTree) -> Self { + type Error = RecoveryError; + + fn try_from(compressed: CompressedMerkleTree) -> Result { + let capacity = 1usize << lb_dynamic_merkle::TREE_HEIGHT_EXCEPT_ROOT; + let mut keys = HashSet::with_capacity(compressed.items.len()); + for (position, (key, _)) in &compressed.items { + if *position >= capacity { + return Err(RecoveryError::PositionOutOfBounds { + position: *position, + capacity, + }); + } + if !keys.insert(key) { + return Err(RecoveryError::DuplicateKey { + position: *position, + }); + } + } + // `items` is a `BTreeMap`, so iteration is ordered by position. let merkle = DynamicMerkleTree::from_sorted_items( compressed @@ -282,7 +312,7 @@ where .iter() .map(|(pos, (key, item))| (*pos, Leaf::leaf(key, item))), ); - Self { + Ok(Self { merkle, items: compressed .items @@ -290,16 +320,61 @@ where .map(|(pos, (key, item))| (key.clone(), (item.clone(), *pos))) .collect(), _leaf: PhantomData, - } + }) } } -#[derive(::serde::Serialize, ::serde::Deserialize)] +#[derive(::serde::Serialize)] #[serde(transparent)] pub struct CompressedMerkleTree { items: BTreeMap, } +// Deserialize manually so duplicate serialized positions are rejected before +// they can be collapsed by `BTreeMap`. A derived implementation would lose +// that information while constructing the map. +impl<'de, Key, Item> ::serde::Deserialize<'de> for CompressedMerkleTree +where + Key: ::serde::Deserialize<'de>, + Item: ::serde::Deserialize<'de>, +{ + fn deserialize(deserializer: D) -> Result + where + D: ::serde::Deserializer<'de>, + { + struct CompressedMerkleTreeVisitor(PhantomData<(Key, Item)>); + + impl<'de, Key, Item> ::serde::de::Visitor<'de> for CompressedMerkleTreeVisitor + where + Key: ::serde::Deserialize<'de>, + Item: ::serde::Deserialize<'de>, + { + type Value = CompressedMerkleTree; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a map of Merkle tree positions to entries") + } + + fn visit_map(self, mut access: A) -> Result + where + A: ::serde::de::MapAccess<'de>, + { + let mut items = BTreeMap::new(); + while let Some((position, item)) = access.next_entry()? { + if items.insert(position, item).is_some() { + return Err(::serde::de::Error::custom( + "compressed Merkle tree contains duplicate positions", + )); + } + } + Ok(CompressedMerkleTree { items }) + } + } + + deserializer.deserialize_map(CompressedMerkleTreeVisitor(PhantomData)) + } +} + mod serde { use serde::{Deserialize, Deserializer, Serialize, Serializer}; @@ -330,7 +405,7 @@ mod serde { D: Deserializer<'de>, { let compressed = super::CompressedMerkleTree::::deserialize(deserializer)?; - Ok(compressed.into()) + Self::try_from(compressed).map_err(serde::de::Error::custom) } } } diff --git a/merkle/utxotree/src/lib.rs b/merkle/utxotree/src/lib.rs index be0de9e27..c3431b341 100644 --- a/merkle/utxotree/src/lib.rs +++ b/merkle/utxotree/src/lib.rs @@ -9,7 +9,7 @@ use ark_ff::AdditiveGroup; pub use lb_dynamic_merkle::{DynamicMerkleTree, MerkleNode, MerklePath, TREE_HEIGHT_EXCEPT_ROOT}; use lb_dynamic_merkle::{MerkleHasher, empty_subtree_root}; pub use lb_merkle_tree::Error; -use lb_merkle_tree::{CompressedMerkleTree, LeafExtractor, MerkleTree}; +use lb_merkle_tree::{CompressedMerkleTree, LeafExtractor, MerkleTree, RecoveryError}; use lb_poseidon2::{Digest, Fr}; use rpds::HashTrieMapSync; @@ -201,14 +201,16 @@ where } } -impl From> for UtxoTree +impl TryFrom> for UtxoTree where Key: AsRef + Clone + std::hash::Hash + Eq, Hash: Digest, Item: Clone, { - fn from(compressed: CompressedUtxoTree) -> Self { - Self(compressed.0.into()) + type Error = RecoveryError; + + fn try_from(compressed: CompressedUtxoTree) -> Result { + Ok(Self(compressed.0.try_into()?)) } } @@ -247,7 +249,7 @@ mod serde { D: Deserializer<'de>, { let compressed = super::CompressedUtxoTree::::deserialize(deserializer)?; - Ok(compressed.into()) + Self::try_from(compressed).map_err(serde::de::Error::custom) } } } @@ -520,7 +522,7 @@ mod tests { let compressed = original_tree.compressed(); // Recover the tree from compressed format - let recovered_tree: UtxoTree<_, _, _> = compressed.into(); + let recovered_tree: UtxoTree<_, _, _> = compressed.try_into().unwrap(); recovered_tree == original_tree }