smt: add from_raw_parts() to trait interface

src/merkle/smt/full/mod.rs

@@ -81,6 +81,14 @@ impl Smt {
         <Self as SparseMerkleTree<SMT_DEPTH>>::with_entries(entries)
     }
 
+    pub fn from_raw_parts(
+        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, SmtLeaf>,
+        root: RpoDigest,
+    ) -> Result<Self, MerkleError> {
+        <Self as SparseMerkleTree<SMT_DEPTH>>::from_raw_parts(inner_nodes, leaves, root)
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -247,6 +255,19 @@ impl SparseMerkleTree<SMT_DEPTH> 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<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, SmtLeaf>,
+        root: RpoDigest,
+    ) -> Result<Self, MerkleError> {
+        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 with_entries(
         entries: impl IntoIterator<Item = (RpoDigest, Word)>,
     ) -> Result<Self, MerkleError> {

src/merkle/smt/mod.rs

@@ -299,6 +299,14 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     where
         Self: Sized;
 
+    fn from_raw_parts(
+        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, Self::Leaf>,
+        root: RpoDigest,
+    ) -> Result<Self, MerkleError>
+    where
+        Self: Sized;
+
     /// The root of the tree
     fn root(&self) -> RpoDigest;
 

src/merkle/smt/simple/mod.rs

@@ -100,6 +100,14 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         Ok(tree)
     }
 
+    pub fn from_raw_parts(
+        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, Word>,
+        root: RpoDigest,
+    ) -> Result<Self, MerkleError> {
+        <Self as SparseMerkleTree<DEPTH>>::from_raw_parts(inner_nodes, leaves, root)
+    }
+
     /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
     /// starting at index 0.
     pub fn with_contiguous_leaves(
@@ -309,6 +317,19 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
     const EMPTY_VALUE: Self::Value = EMPTY_WORD;
     const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
 
+    fn from_raw_parts(
+        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, Word>,
+        root: RpoDigest,
+    ) -> Result<Self, MerkleError> {
+        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 with_entries(
         entries: impl IntoIterator<Item = (LeafIndex<DEPTH>, Word)>,
     ) -> Result<Self, MerkleError> {