diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index fd9a4f59..4f6ec622 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -101,6 +101,23 @@ impl Smt { Ok(tree) } + /// Returns a new [`Smt`] instantiated from already computed leaves and nodes. + /// + /// This function performs minimal consistency checking. It is the caller's responsibility to + /// ensure the passed arguments are correct and consistent with each other. + /// + /// # Panics + /// With debug assertions on, this function panics if `root` does not match the root node in + /// `inner_nodes`. + pub fn from_raw_parts( + inner_nodes: BTreeMap, + leaves: BTreeMap, + root: RpoDigest, + ) -> Self { + // Our particular implementation of `from_raw_parts()` never returns `Err`. + >::from_raw_parts(inner_nodes, leaves, root).unwrap() + } + // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- @@ -280,6 +297,19 @@ impl SparseMerkleTree for Smt { const EMPTY_VALUE: Self::Value = EMPTY_WORD; const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); + fn from_raw_parts( + inner_nodes: BTreeMap, + leaves: BTreeMap, + root: RpoDigest, + ) -> Result { + if cfg!(debug_assertions) { + let root_node = inner_nodes.get(&NodeIndex::root()).unwrap(); + assert_eq!(root_node.hash(), root); + } + + Ok(Self { root, inner_nodes, leaves }) + } + fn root(&self) -> RpoDigest { self.root } diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index a9ed32c5..adeecd9d 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -292,6 +292,16 @@ pub(crate) trait SparseMerkleTree { // REQUIRED METHODS // --------------------------------------------------------------------------------------------- + /// Construct this type from already computed leaves and nodes. The caller ensures passed + /// arguments are correct and consistent with each other. + fn from_raw_parts( + inner_nodes: BTreeMap, + leaves: BTreeMap, + root: RpoDigest, + ) -> Result + where + Self: Sized; + /// The root of the tree fn root(&self) -> RpoDigest; diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 04476a0a..1ded87f1 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -100,6 +100,23 @@ impl SimpleSmt { Ok(tree) } + /// Returns a new [`SimpleSmt`] instantiated from already computed leaves and nodes. + /// + /// This function performs minimal consistency checking. It is the caller's responsibility to + /// ensure the passed arguments are correct and consistent with each other. + /// + /// # Panics + /// With debug assertions on, this function panics if `root` does not match the root node in + /// `inner_nodes`. + pub fn from_raw_parts( + inner_nodes: BTreeMap, + leaves: BTreeMap, + root: RpoDigest, + ) -> Self { + // Our particular implementation of `from_raw_parts()` never returns `Err`. + >::from_raw_parts(inner_nodes, leaves, root).unwrap() + } + /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices /// starting at index 0. pub fn with_contiguous_leaves( @@ -309,6 +326,19 @@ impl SparseMerkleTree for SimpleSmt { const EMPTY_VALUE: Self::Value = EMPTY_WORD; const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0); + fn from_raw_parts( + inner_nodes: BTreeMap, + leaves: BTreeMap, + root: RpoDigest, + ) -> Result { + if cfg!(debug_assertions) { + let root_node = inner_nodes.get(&NodeIndex::root()).unwrap(); + assert_eq!(root_node.hash(), root); + } + + Ok(Self { root, inner_nodes, leaves }) + } + fn root(&self) -> RpoDigest { self.root }