From 1aa09e735da0a88096c044a58eae5d61ba13d0ae Mon Sep 17 00:00:00 2001 From: Daniel Sanchez Date: Fri, 24 Jul 2026 14:59:42 +0000 Subject: [PATCH] feat(merkle): Dynamic merkle tree (#3164) Co-authored-by: thomaslavaur <82142582+thomaslavaur@users.noreply.github.com> --- Cargo.lock | 14 +- Cargo.toml | 2 + dynamic-merkle/Cargo.toml | 23 + .../merkle.rs => dynamic-merkle/src/lib.rs | 498 ++++++++++++------ tests/benches/voucher.rs | 4 +- utxotree/Cargo.toml | 14 +- utxotree/src/lib.rs | 56 +- 7 files changed, 420 insertions(+), 191 deletions(-) create mode 100644 dynamic-merkle/Cargo.toml rename utxotree/src/merkle.rs => dynamic-merkle/src/lib.rs (57%) diff --git a/Cargo.lock b/Cargo.lock index 910d1b9d2..0f9cb2ef3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4567,6 +4567,18 @@ dependencies = [ "tracing-subscriber 0.3.23", ] +[[package]] +name = "logos-blockchain-dynamic-merkle" +version = "0.0.0" +dependencies = [ + "ark-ff", + "logos-blockchain-poseidon2", + "num-bigint", + "rand 0.8.6", + "rpds", + "serde", +] + [[package]] name = "logos-blockchain-faucet" version = "0.0.0" @@ -5252,7 +5264,7 @@ name = "logos-blockchain-utxotree" version = "0.0.0" dependencies = [ "ark-ff", - "logos-blockchain-groth16", + "logos-blockchain-dynamic-merkle", "logos-blockchain-poseidon2", "num-bigint", "quickcheck", diff --git a/Cargo.toml b/Cargo.toml index 99d825ba4..8dfc0a9d0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,7 @@ members = [ "deployment/l2-sequencer-archival-demo/archiver", "deployment/l2-sequencer-archival-demo/sequencer", "deployment/tui-zone", + "dynamic-merkle", "kms/keys", "kms/macros", "kms/operators", @@ -125,6 +126,7 @@ lb-cryptarchia-engine = { default-features = false, package = "logo lb-cryptarchia-sync = { default-features = false, package = "logos-blockchain-cryptarchia-sync", path = "./consensus/cryptarchia-sync" } lb-demo-archiver = { default-features = false, package = "logos-blockchain-demo-archiver", path = "./deployment/l2-sequencer-archival-demo/archiver" } lb-demo-sequencer = { default-features = false, package = "logos-blockchain-demo-sequencer", path = "./deployment/l2-sequencer-archival-demo/sequencer" } +lb-dynamic-merkle = { default-features = false, package = "logos-blockchain-dynamic-merkle", path = "./dynamic-merkle" } lb-faucet = { default-features = false, package = "logos-blockchain-faucet", path = "./deployment/faucet" } lb-groth16 = { default-features = false, package = "logos-blockchain-groth16", path = "./zk/groth16" } lb-http-api-common = { default-features = false, package = "logos-blockchain-http-api-common", path = "./nodes/api-common" } diff --git a/dynamic-merkle/Cargo.toml b/dynamic-merkle/Cargo.toml new file mode 100644 index 000000000..ced714fe5 --- /dev/null +++ b/dynamic-merkle/Cargo.toml @@ -0,0 +1,23 @@ +[package] +categories.workspace = true +description.workspace = true +edition.workspace = true +keywords.workspace = true +license.workspace = true +name = "logos-blockchain-dynamic-merkle" +readme.workspace = true +repository.workspace = true +version.workspace = true + +[dependencies] +rpds = { features = ["serde"], workspace = true } +serde = { features = ["alloc", "derive", "rc"], workspace = true } + +[dev-dependencies] +ark-ff = { workspace = true } +lb-poseidon2 = { workspace = true } +num-bigint = { workspace = true } +rand = { workspace = true } + +[lints] +workspace = true diff --git a/utxotree/src/merkle.rs b/dynamic-merkle/src/lib.rs similarity index 57% rename from utxotree/src/merkle.rs rename to dynamic-merkle/src/lib.rs index 26f9a800e..280c34e50 100644 --- a/utxotree/src/merkle.rs +++ b/dynamic-merkle/src/lib.rs @@ -1,41 +1,98 @@ -use std::{ - marker::PhantomData, - sync::{Arc, OnceLock}, -}; +//! A dynamic, persistent, fixed-height Merkle tree generic over its hashing +//! backend. +//! +//! The tree stores items at leaf positions of a binary tree of fixed height +//! ([`TREE_HEIGHT_EXCEPT_ROOT`]). Insertions fill the lowest available position +//! (reusing positions freed by removals), so an item's index is stable for the +//! lifetime of the tree and membership proofs have a constant length. +//! +//! The item type, value/hash type and hashing operations are supplied by a +//! [`MerkleHasher`] implementation, which the tree is parameterized over. Use +//! the [`empty_subtree_root`] macro to derive the cached +//! [`MerkleHasher::empty_subtree_root`] method for a concrete hash type. + +use std::{fmt, marker::PhantomData, sync::Arc}; -use ark_ff::AdditiveGroup; -use lb_groth16::serde::serde_fr; -use lb_poseidon2::{Digest, Fr}; use rpds::RedBlackTreeSetSync; -use crate::CompressedUtxoTree; +/// Abstraction over the item/hash types and hashing operations a +/// [`DynamicMerkleTree`] needs. +/// +/// A single implementor binds together the leaf payload type ([`Self::Item`]) +/// and the field/value type stored in inner nodes, roots and paths +/// ([`Self::Hash`]). +pub trait MerkleHasher { + /// The leaf payload stored in the tree. + type Item: Clone; + /// The value type: inner node values, roots and merkle-path siblings. + type Hash: Copy + Eq; -const TREE_HEIGHT_EXCEPT_ROOT: usize = 32; + /// Neutral value used for empty leaves and as the seed of empty subtrees. + const EMPTY_VALUE: Self::Hash; -const EMPTY_VALUE: Fr = ::ZERO; + /// Extract the hash value of a leaf item. + fn leaf_hash(item: &Self::Item) -> Self::Hash; -fn empty_subtree_root(height: usize) -> Fr { - static PRECOMPUTED_EMPTY_ROOTS: OnceLock<[Fr; TREE_HEIGHT_EXCEPT_ROOT + 1]> = OnceLock::new(); - assert!( - height <= TREE_HEIGHT_EXCEPT_ROOT, - "Height{height} must be <={TREE_HEIGHT_EXCEPT_ROOT}" - ); - PRECOMPUTED_EMPTY_ROOTS.get_or_init(|| { - let mut hashes = [EMPTY_VALUE; TREE_HEIGHT_EXCEPT_ROOT + 1]; - for i in 1..=TREE_HEIGHT_EXCEPT_ROOT { - hashes[i] = Hash::compress(&[hashes[i - 1], hashes[i - 1]]); + /// Compress two child hashes into their parent hash. + fn compress(left: &Self::Hash, right: &Self::Hash) -> Self::Hash; + + /// Root of a fully-empty subtree of the given `height`. + /// + /// Implement with [`empty_subtree_root`] to get a cached implementation. + fn empty_subtree_root(height: usize) -> Self::Hash; +} + +/// Height of the tree excluding the root, i.e. the length of every Merkle path +/// and the base-2 logarithm of the tree's leaf capacity (`2^32` items). +pub const TREE_HEIGHT_EXCEPT_ROOT: usize = 32; + +/// Generates a cached [`MerkleHasher::empty_subtree_root`] implementation for a +/// concrete `Hash` type. +/// +/// The cache is a `static` local to the generated method, so it is +/// monomorphization-free (the `Hash` type is concrete here) and each +/// implementing type gets its own independent cache. +/// +/// ```ignore +/// impl MerkleHasher for MyHasher { +/// type Item = MyItem; +/// type Hash = Fr; +/// const EMPTY_VALUE: Fr = /* ... */; +/// fn leaf_hash(item: &MyItem) -> Fr { /* ... */ } +/// fn compress(left: &Fr, right: &Fr) -> Fr { /* ... */ } +/// empty_subtree_root!(Fr); +/// } +/// ``` +#[macro_export] +macro_rules! empty_subtree_root { + ($hash:ty) => { + fn empty_subtree_root(height: usize) -> $hash { + static PRECOMPUTED_EMPTY_ROOTS: ::std::sync::OnceLock< + [$hash; $crate::TREE_HEIGHT_EXCEPT_ROOT + 1], + > = ::std::sync::OnceLock::new(); + assert!( + height <= $crate::TREE_HEIGHT_EXCEPT_ROOT, + "Height{height} must be <={}", + $crate::TREE_HEIGHT_EXCEPT_ROOT + ); + PRECOMPUTED_EMPTY_ROOTS.get_or_init(|| { + let mut hashes = [Self::EMPTY_VALUE; $crate::TREE_HEIGHT_EXCEPT_ROOT + 1]; + for i in 1..=$crate::TREE_HEIGHT_EXCEPT_ROOT { + hashes[i] = Self::compress(&hashes[i - 1], &hashes[i - 1]); + } + hashes + })[height] } - hashes - })[height] + }; } #[derive(::serde::Serialize, ::serde::Deserialize, Clone, Debug, PartialEq, Eq)] -enum Node { +enum Node { Inner { left: Arc, right: Arc, - #[serde(with = "serde_fr")] - value: Fr, + // Hash is bound to a value, not to confuse with Hasher + value: Hash, right_subtree_size: usize, left_subtree_size: usize, height: usize, @@ -52,28 +109,21 @@ enum Node { }, } -fn hash, Hash: Digest>(left: &Node, right: &Node) -> Fr { - let mut input = [EMPTY_VALUE; 2]; - match left { - Node::Inner { value, .. } => input[0] = *value, - Node::Leaf { item } => { - input[0] = *item.as_ref().map_or(&EMPTY_VALUE, AsRef::as_ref); - } +fn hash(left: &Node, right: &Node) -> H::Hash { + let left = match left { + Node::Inner { value, .. } => *value, + Node::Leaf { item } => item.as_ref().map_or(H::EMPTY_VALUE, H::leaf_hash), Node::Empty { .. } => panic!("Empty node in left subtree is not allowed"), - } - match right { - Node::Inner { value, .. } => input[1] = *value, - Node::Leaf { item } => { - input[1] = *item.as_ref().map_or(&EMPTY_VALUE, AsRef::as_ref); - } - Node::Empty { height } => { - input[1] = empty_subtree_root::(*height); - } - } - Hash::compress(&input) + }; + let right = match right { + Node::Inner { value, .. } => *value, + Node::Leaf { item } => item.as_ref().map_or(H::EMPTY_VALUE, H::leaf_hash), + Node::Empty { height } => H::empty_subtree_root(*height), + }; + H::compress(&left, &right) } -impl Node { +impl Node { const fn new(item: Item) -> Self { Self::Leaf { item: Some(item) } } @@ -103,28 +153,28 @@ impl Node { } } -impl> Node { - fn new_inner(left: Arc, right: Arc) -> Self +impl Node { + fn new_inner(left: Arc, right: Arc) -> Self where - Hash: Digest, + H: MerkleHasher, { Self::Inner { right_subtree_size: right.size(), left_subtree_size: left.size(), height: left.height().max(right.height()) + 1, - value: hash::<_, Hash>(&left, &right), + value: hash::(&left, &right), left, right, } } - fn insert_or_modify Self>( + fn insert_or_modify Self>( self: &Arc, index: usize, f: F, ) -> Arc where - Hash: Digest, + H: MerkleHasher, { match self.as_ref() { Self::Inner { left, right, .. } => { @@ -137,15 +187,15 @@ impl> Node { if index < left.capacity() { // modify the left subtree - Arc::new(Self::new_inner::( - left.insert_or_modify::(index, f), + Arc::new(Self::new_inner::( + left.insert_or_modify::(index, f), Arc::clone(right), )) } else { // modify the right subtree - Arc::new(Self::new_inner::( + Arc::new(Self::new_inner::( Arc::clone(left), - right.insert_or_modify::(index - left.capacity(), f), + right.insert_or_modify::(index - left.capacity(), f), )) } } @@ -155,9 +205,8 @@ impl> Node { index == 0, "Cannot expand an empty subtree more than one node at a time", ); - Arc::new(Self::new_inner::( - Arc::new(Self::Empty { height: height - 1 }) - .insert_or_modify::(index, f), + Arc::new(Self::new_inner::( + Arc::new(Self::Empty { height: height - 1 }).insert_or_modify::(index, f), Arc::new(Self::Empty { height: height - 1 }), )) } @@ -171,22 +220,22 @@ impl> Node { } } - fn insert_at(self: &Arc, index: usize, item: Item) -> Arc + fn insert_at(self: &Arc, index: usize, item: Item) -> Arc where - Hash: Digest, + H: MerkleHasher, { - self.insert_or_modify::(index, |node| match node { + self.insert_or_modify::(index, |node| match node { Self::Leaf { item: None } | Self::Empty { .. } => Self::new(item), Self::Leaf { item: Some(_) } => panic!("Cannot insert into a non-empty leaf node"), _ => panic!("Cannot insert into a non-terminal node"), }) } - fn remove_at(self: &Arc, index: usize) -> Arc + fn remove_at(self: &Arc, index: usize) -> Arc where - Hash: Digest, + H: MerkleHasher, { - self.insert_or_modify::(index, move |node| match node { + self.insert_or_modify::(index, move |node| match node { Self::Leaf { item: Some(_) } => Self::Leaf { item: None }, _ => panic!("Cannot remove from a empty / non-leaf node"), }) @@ -195,9 +244,9 @@ impl> Node { /// Computes the Merkle path for the item at the given index. /// The path is ordered from leaf to root (excluded). /// Returns `None` if the index does not exist or has been removed. - fn path(self: &Arc, index: usize) -> Option> + fn path(self: &Arc, index: usize) -> Option> where - Hash: Digest, + H: MerkleHasher, { match self.as_ref() { Self::Inner { left, right, .. } => { @@ -210,15 +259,15 @@ impl> Node { if index < left.capacity() { // Going down left subtree, store right sibling hash - let mut path = left.path::(index)?; + let mut path = left.path::(index)?; assert!(path.len() < TREE_HEIGHT_EXCEPT_ROOT, "Path length exceeded"); - path.push(MerkleNode::Right(right.value::())); + path.push(MerkleNode::Right(right.value::())); Some(path) } else { // Going down right subtree, store left sibling hash - let mut path = right.path::(index - left.capacity())?; + let mut path = right.path::(index - left.capacity())?; assert!(path.len() < TREE_HEIGHT_EXCEPT_ROOT, "Path length exceeded"); - path.push(MerkleNode::Left(left.value::())); + path.push(MerkleNode::Left(left.value::())); Some(path) } } @@ -227,15 +276,15 @@ impl> Node { } } - fn value(&self) -> Fr + fn value(&self) -> Hash where - Hash: Digest, + H: MerkleHasher, { match self { Self::Inner { value, .. } => *value, - Self::Leaf { item: Some(item) } => *item.as_ref(), - Self::Leaf { item: None } => EMPTY_VALUE, - Self::Empty { height } => empty_subtree_root::(*height), + Self::Leaf { item: Some(item) } => H::leaf_hash(item), + Self::Leaf { item: None } => H::EMPTY_VALUE, + Self::Empty { height } => H::empty_subtree_root(*height), } } } @@ -246,16 +295,38 @@ impl> Node { /// Removed items 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, allowing for efficient and simple proof of -/// memberships for `PoL`. -#[derive(Debug, Clone)] -pub struct DynamicMerkleTree { - root: Arc>, +/// 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, - _hash: PhantomData, + _hasher: PhantomData, } -impl, Hash: Digest> Default for DynamicMerkleTree { +impl Clone for DynamicMerkleTree { + fn clone(&self) -> Self { + Self { + root: Arc::clone(&self.root), + holes: self.holes.clone(), + _hasher: PhantomData, + } + } +} + +impl fmt::Debug for DynamicMerkleTree +where + H::Item: fmt::Debug, + H::Hash: fmt::Debug, +{ + 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 { @@ -263,23 +334,38 @@ impl, Hash: Digest> Default for DynamicMerkleTree { height: TREE_HEIGHT_EXCEPT_ROOT, }), holes, - _hash: PhantomData, + _hasher: PhantomData, } } } -impl, Hash: Digest> DynamicMerkleTree { +impl DynamicMerkleTree { + /// Creates a new, empty tree. #[must_use] pub fn new() -> Self { Self::default() } + /// Returns the number of items currently stored in the tree (removed + /// positions do not count). #[must_use] pub fn size(&self) -> usize { self.root.size() } - pub fn insert(&self, item: Item) -> (Self, usize) { + /// Inserts `item` at the lowest available position and returns the updated + /// tree together with the index the item was assigned. + /// + /// Positions freed by [`remove`](Self::remove) are reused before the tree + /// grows, so the smallest free index is always chosen. + /// + /// The original tree is left unchanged (the structure is persistent). + /// + /// # Panics + /// + /// Panics if the tree is already at full capacity + /// (`2^TREE_HEIGHT_EXCEPT_ROOT` items). + pub fn insert(&self, item: H::Item) -> (Self, usize) { assert!( self.size() < self.root.capacity(), "max capacity reached, cannot insert more items" @@ -290,37 +376,51 @@ impl, Hash: Digest> DynamicMerkleTree { |hole| (self.holes.remove(hole), *hole), ); - let root = self.root.insert_at::(index, item); + let root = self.root.insert_at::(index, item); ( Self { root, holes, - _hash: PhantomData, + _hasher: PhantomData, }, index, ) } - pub(crate) fn remove(&self, index: usize) -> Self { + /// Removes the item 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. + /// + /// # Panics + /// + /// Panics if `index` is out of bounds, or if the position does not hold an + /// item. + #[must_use] + pub fn remove(&self, index: usize) -> Self { assert!(index < self.root.capacity(), "Index out of bounds"); - let root = self.root.remove_at::(index); + let root = self.root.remove_at::(index); let holes = self.holes.insert(index); Self { root, holes, - _hash: PhantomData, + _hasher: PhantomData, } } + /// Returns the Merkle root of the tree. + /// + /// An empty tree yields the empty-subtree root for the full height. #[must_use] - pub fn root(&self) -> Fr { + pub fn root(&self) -> H::Hash { match self.root.as_ref() { Node::Inner { value, .. } => *value, Node::Leaf { .. } => { panic!("Cannot get root from a leaf node, expected an inner node or empty node"); } - Node::Empty { .. } => empty_subtree_root::(self.root.height()), + Node::Empty { .. } => H::empty_subtree_root(self.root.height()), } } @@ -328,8 +428,8 @@ impl, Hash: Digest> DynamicMerkleTree { /// The path is ordered from leaf to root (excluded). /// Returns `None` if the index does not exist or has been removed. #[must_use] - pub fn path(&self, index: usize) -> Option> { - self.root.path::(index).inspect(|path| { + pub fn path(&self, index: usize) -> Option> { + self.root.path::(index).inspect(|path| { assert_eq!( path.len(), TREE_HEIGHT_EXCEPT_ROOT, @@ -339,6 +439,34 @@ impl, Hash: Digest> DynamicMerkleTree { }) } + /// Rebuilds a tree placing each `item` at its given index, filling the gaps + /// between indices with holes. + /// + /// The items must be yielded in strictly increasing index order; this is + /// the inverse of enumerating a tree's occupied positions and is meant + /// for recovering a tree from a compressed representation. + /// + /// # Panics + /// + /// Panics if the indices are not strictly increasing or an index is out of + /// bounds. + #[must_use] + pub fn from_sorted_items(items: impl IntoIterator) -> Self { + let mut tree = Self::new(); + let mut current_pos = 0; + for (pos, item) 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, item); + 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 { @@ -350,7 +478,7 @@ impl, Hash: Digest> DynamicMerkleTree { let holes = self.holes.insert(index); let root = self .root - .insert_or_modify::(index, |node| match node { + .insert_or_modify::(index, |node| match node { Node::Empty { .. } => Node::Leaf { item: None }, _ => panic!("Cannot insert a hole into a non-empty/non-leaf node"), }); @@ -358,71 +486,44 @@ impl, Hash: Digest> DynamicMerkleTree { Self { root, holes, - _hash: PhantomData, + _hasher: PhantomData, } } } -impl + Clone, Hash: Digest> DynamicMerkleTree { - pub(crate) fn from_compressed_tree(comp: &CompressedUtxoTree) -> Self { - let mut tree = Self::new(); - let mut current_pos = 0; - for (pos, (key, _)) in &comp.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, key.clone()); - current_pos = *pos + 1; - } - tree - } -} - -impl PartialEq for DynamicMerkleTree -where - Item: AsRef + PartialEq, - Hash: Digest, -{ +impl PartialEq for DynamicMerkleTree { fn eq(&self, other: &Self) -> bool { self.root() == other.root() } } -impl Eq for DynamicMerkleTree -where - Item: AsRef + Eq, - Hash: Digest, -{ -} +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 +/// [`Item`](MerkleHasher::Item) and [`Hash`](MerkleHasher::Hash) types to +/// implement the corresponding `serde` traits. pub mod serde { use std::{marker::PhantomData, sync::Arc}; use rpds::RedBlackTreeSetSync; use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeStruct as _}; + use super::MerkleHasher; + #[derive(Deserialize)] - pub struct DynamicMerkleTree { - root: Arc>, + struct Raw { + root: Arc>, holes: RedBlackTreeSetSync, } - impl From> for super::DynamicMerkleTree { - fn from(tree: DynamicMerkleTree) -> Self { - Self { - root: tree.root, - holes: tree.holes, - _hash: PhantomData, - } - } - } - - impl Serialize for super::DynamicMerkleTree + impl Serialize for super::DynamicMerkleTree where - Item: Serialize, + H: MerkleHasher, + H::Item: Serialize, + H::Hash: Serialize, { fn serialize(&self, serializer: S) -> Result where @@ -435,16 +536,22 @@ pub mod serde { } } - impl<'de, Item, Hash> Deserialize<'de> for super::DynamicMerkleTree + impl<'de, H> Deserialize<'de> for super::DynamicMerkleTree where - Item: Deserialize<'de>, + H: MerkleHasher, + H::Item: Deserialize<'de>, + H::Hash: Deserialize<'de>, { fn deserialize(deserializer: D) -> Result where D: Deserializer<'de>, { - let raw: DynamicMerkleTree = Deserialize::deserialize(deserializer)?; - Ok(raw.into()) + let raw = Raw::::deserialize(deserializer)?; + Ok(Self { + root: raw.root, + holes: raw.holes, + _hasher: PhantomData, + }) } } } @@ -459,6 +566,7 @@ pub enum MerkleNode { } impl MerkleNode { + /// Returns the sibling value, regardless of which side it is on. pub const fn item(&self) -> &T { match self { Self::Left(v) | Self::Right(v) => v, @@ -470,26 +578,78 @@ impl MerkleNode { pub type MerklePath = Vec>; #[cfg(test)] -mod tests { - use super::*; - use crate::test_fr::TestFr; +mod test_fr { + use ark_ff::AdditiveGroup; + use lb_poseidon2::{Digest, Fr, Poseidon2Bn254Hasher}; + use num_bigint::BigUint; + use rand::RngCore; - type TestHash = lb_poseidon2::Poseidon2Bn254Hasher; + use crate::MerkleHasher; + + #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] + pub struct TestFr(Fr); + + impl TestFr { + pub fn from_rng(rng: &mut Rng) -> Self { + Self(BigUint::from(rng.next_u64()).into()) + } + + #[must_use] + pub fn from_usize(n: usize) -> Self { + Self(BigUint::from(n).into()) + } + } + + impl AsRef for TestFr { + fn as_ref(&self) -> &Fr { + &self.0 + } + } + + /// Test [`MerkleHasher`] backed by Poseidon2 over BN254. + pub struct TestHasher; + + impl MerkleHasher for TestHasher { + type Item = TestFr; + type Hash = Fr; + + const EMPTY_VALUE: Fr = ::ZERO; + + fn leaf_hash(item: &TestFr) -> Fr { + *item.as_ref() + } + + fn compress(left: &Fr, right: &Fr) -> Fr { + ::compress(&[*left, *right]) + } + + empty_subtree_root!(Fr); + } +} + +#[cfg(test)] +mod tests { + use lb_poseidon2::Fr; + + use super::{ + test_fr::{TestFr, TestHasher}, + *, + }; #[test] fn test_empty_tree() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); assert_eq!(tree.size(), 0); assert_eq!( tree.root(), - empty_subtree_root::(TREE_HEIGHT_EXCEPT_ROOT) + TestHasher::empty_subtree_root(TREE_HEIGHT_EXCEPT_ROOT) ); assert_eq!(tree.root.height(), TREE_HEIGHT_EXCEPT_ROOT); } #[test] fn test_hole_management() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let mut rng = rand::thread_rng(); let a = TestFr::from_rng(&mut rng); let b = TestFr::from_rng(&mut rng); @@ -509,7 +669,7 @@ mod tests { #[test] fn test_root_consistency() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let mut rng = rand::thread_rng(); let a = TestFr::from_rng(&mut rng); let b = TestFr::from_rng(&mut rng); @@ -530,11 +690,11 @@ mod tests { let mut rng = rand::thread_rng(); let a = TestFr::from_rng(&mut rng); let b = TestFr::from_rng(&mut rng); - let tree1: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree1: DynamicMerkleTree = DynamicMerkleTree::new(); let (tree1, _) = tree1.insert(a); let (tree1, _) = tree1.insert(b); - let tree2: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree2: DynamicMerkleTree = DynamicMerkleTree::new(); let (tree2, _) = tree2.insert(a); let (tree2, _) = tree2.insert(b); @@ -544,14 +704,14 @@ mod tests { #[test] #[should_panic(expected = "Index out of bounds")] fn test_remove_out_of_bounds() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let (tree, _) = tree.insert(TestFr::from_rng(&mut rand::thread_rng())); - tree.remove(1 << 32); + let _tree = tree.remove(1 << 32); } #[test] fn test_single_insert() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let item = TestFr::from_rng(&mut rand::thread_rng()); let (tree_with_item, index) = tree.insert(item); @@ -563,7 +723,7 @@ mod tests { #[test] fn test_multiple_inserts() { - let mut tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let mut tree: DynamicMerkleTree = DynamicMerkleTree::new(); let items = [ TestFr::from_rng(&mut rand::thread_rng()), TestFr::from_rng(&mut rand::thread_rng()), @@ -582,7 +742,7 @@ mod tests { #[test] fn test_remove_single_item() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); let item = TestFr::from_rng(&mut rand::thread_rng()); let (tree_with_item, _) = tree.insert(item); @@ -593,7 +753,7 @@ mod tests { #[test] fn test_remove_and_reinsert() { - let mut tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let mut tree: DynamicMerkleTree = DynamicMerkleTree::new(); let items = vec![ TestFr::from_rng(&mut rand::thread_rng()), TestFr::from_rng(&mut rand::thread_rng()), @@ -616,7 +776,7 @@ mod tests { #[test] fn test_structural_sharing() { - let tree1: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree1: DynamicMerkleTree = DynamicMerkleTree::new(); let (tree2, _) = tree1.insert(TestFr::from_rng(&mut rand::thread_rng())); let (tree3, _) = tree2.insert(TestFr::from_rng(&mut rand::thread_rng())); @@ -631,7 +791,7 @@ mod tests { #[test] fn test_smallest_hole_selection() { - let tree: DynamicMerkleTree = DynamicMerkleTree::new(); + let tree: DynamicMerkleTree = DynamicMerkleTree::new(); // Insert items at positions 0, 1, 2, 3, 4 let (tree, _) = tree.insert(TestFr::from_rng(&mut rand::thread_rng())); @@ -661,7 +821,7 @@ mod tests { #[test] fn test_path_empty_tree() { - let tree = DynamicMerkleTree::::new(); + let tree = DynamicMerkleTree::::new(); // Getting a path from an empty tree should return None assert!(tree.path(0).is_none()); @@ -669,7 +829,7 @@ mod tests { #[test] fn test_path_single_item() { - let tree = DynamicMerkleTree::::new(); + let tree = DynamicMerkleTree::::new(); let item = TestFr::from_usize(0); let (tree, idx) = tree.insert(item); @@ -683,14 +843,14 @@ mod tests { // So all siblings should be Right nodes with empty subtree hashes for (height, node) in path.iter().enumerate() { assert!(matches!(node, MerkleNode::Right(_))); - let sibling_hash = empty_subtree_root::(height); + let sibling_hash = TestHasher::empty_subtree_root(height); assert_eq!(*node.item(), sibling_hash); } } #[test] fn test_path_removed_item() { - let tree = DynamicMerkleTree::::new(); + let tree = DynamicMerkleTree::::new(); let (tree, idx) = tree.insert(TestFr::from_usize(0)); // Path should exist before removal @@ -704,7 +864,7 @@ mod tests { #[test] fn test_path_multiple_items() { - let tree = DynamicMerkleTree::::new(); + let tree = DynamicMerkleTree::::new(); let item0 = TestFr::from_usize(0); let item1 = TestFr::from_usize(1); let item2 = TestFr::from_usize(2); @@ -737,12 +897,8 @@ mod tests { let mut current_hash = *item.as_ref(); for node in path { current_hash = match node { - MerkleNode::Left(sibling) => { - ::compress(&[*sibling, current_hash]) - } - MerkleNode::Right(sibling) => { - ::compress(&[current_hash, *sibling]) - } + MerkleNode::Left(sibling) => TestHasher::compress(sibling, ¤t_hash), + MerkleNode::Right(sibling) => TestHasher::compress(¤t_hash, sibling), }; } assert_eq!( diff --git a/tests/benches/voucher.rs b/tests/benches/voucher.rs index 6dd7c5018..81a131dc5 100644 --- a/tests/benches/voucher.rs +++ b/tests/benches/voucher.rs @@ -20,7 +20,7 @@ use lb_core::{ }; use lb_groth16::Fr; use lb_mmr::MerkleMountainRange; -use lb_utxotree::DynamicMerkleTree; +use lb_utxotree::{DynamicMerkleTree, UtxoMerkleHasher}; const SAMPLE_COUNT: u32 = 3; @@ -57,7 +57,7 @@ fn tree_11th_epoch(bencher: Bencher) { } type Mmr = MerkleMountainRange; -type Tree = DynamicMerkleTree; +type Tree = DynamicMerkleTree>; fn voucher(i: u64) -> VoucherCm { VoucherCm::from_secret(VoucherSecret::from(Fr::from(i))) diff --git a/utxotree/Cargo.toml b/utxotree/Cargo.toml index 381b32064..5ea874310 100644 --- a/utxotree/Cargo.toml +++ b/utxotree/Cargo.toml @@ -10,15 +10,15 @@ repository = { workspace = true } version = { workspace = true } [dependencies] -ark-ff = { workspace = true } -lb-groth16 = { workspace = true } -lb-poseidon2 = { workspace = true } -num-bigint = { workspace = true } -rpds = { features = ["serde"], workspace = true } -serde = { features = ["alloc", "derive", "rc"], workspace = true } -thiserror = { workspace = true } +ark-ff = { workspace = true } +lb-dynamic-merkle = { workspace = true } +lb-poseidon2 = { workspace = true } +rpds = { workspace = true } +serde = { features = ["derive"], workspace = true } +thiserror = { workspace = true } [dev-dependencies] +num-bigint = { workspace = true } quickcheck = { workspace = true } quickcheck_macros = { workspace = true } rand = { workspace = true } diff --git a/utxotree/src/lib.rs b/utxotree/src/lib.rs index e2ea31ad1..b6f2b13be 100644 --- a/utxotree/src/lib.rs +++ b/utxotree/src/lib.rs @@ -1,17 +1,45 @@ -mod merkle; - #[cfg(test)] pub mod test_fr; -use std::collections::BTreeMap; +use std::{collections::BTreeMap, marker::PhantomData}; -use lb_poseidon2::{Digest, Fr}; +use ark_ff::AdditiveGroup; // TODO: Change `DynamicMerkleTree` back to private once we adopt MMR for vouchers in the // wallet service. -pub use merkle::{DynamicMerkleTree, MerkleNode, MerklePath}; +pub use lb_dynamic_merkle::{DynamicMerkleTree, MerkleNode, MerklePath}; +use lb_dynamic_merkle::{MerkleHasher, empty_subtree_root}; +use lb_poseidon2::{Digest, Fr}; use rpds::HashTrieMapSync; use thiserror::Error; +/// [`MerkleHasher`] bridge adapting a `Key: AsRef` leaf type and a +/// [`Digest`] hasher to the generic [`DynamicMerkleTree`]. +/// +/// Leaf values are the key's field element and inner nodes are compressed with +/// `Hash`. +pub struct UtxoMerkleHasher(PhantomData<(Key, Hash)>); + +impl MerkleHasher for UtxoMerkleHasher +where + Key: AsRef + Clone, + Hash: Digest, +{ + type Item = Key; + type Hash = Fr; + + const EMPTY_VALUE: Fr = ::ZERO; + + fn leaf_hash(item: &Key) -> Fr { + *item.as_ref() + } + + fn compress(left: &Fr, right: &Fr) -> Fr { + ::compress(&[*left, *right]) + } + + empty_subtree_root!(Fr); +} + /// A store for `UTxOs` that allows for efficient insertion, removal, and /// retrieval of items, while efficiently maintaining a compact Merkle tree /// for Proof of Leadership (`PoL`) generation. @@ -24,9 +52,10 @@ use thiserror::Error; #[derive(Debug, Clone)] pub struct UtxoTree where - Key: std::hash::Hash + Eq, + Key: AsRef + Clone + std::hash::Hash + Eq, + Hash: Digest, { - merkle: DynamicMerkleTree, + merkle: DynamicMerkleTree>, // key -> (item, position in merkle tree) items: HashTrieMapSync, } @@ -136,7 +165,7 @@ where impl PartialEq for UtxoTree where - Key: AsRef + std::hash::Hash + Eq, + Key: AsRef + Clone + std::hash::Hash + Eq, Item: PartialEq, Hash: Digest, { @@ -147,7 +176,7 @@ where impl Eq for UtxoTree where - Key: AsRef + std::hash::Hash + Eq, + Key: AsRef + Clone + std::hash::Hash + Eq, Item: Eq, Hash: Digest, { @@ -176,8 +205,15 @@ where Item: Clone, { fn from(compressed: CompressedUtxoTree) -> Self { + // `items` is a `BTreeMap`, so iteration is ordered by position. + let merkle = DynamicMerkleTree::from_sorted_items( + compressed + .items + .iter() + .map(|(pos, (key, _))| (*pos, key.clone())), + ); Self { - merkle: DynamicMerkleTree::from_compressed_tree(&compressed), + merkle, items: compressed .items .iter()