smt: add from_raw_parts() to trait interface

This commit adds a new required method to the SparseMerkleTree trait,
to allow generic construction from pre-computed parts.

This will be used to add a generic version of `with_entries()` in a
later commit.

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<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, SmtLeaf>,
+        root: RpoDigest,
+    ) -> Self {
+        // Our particular implementation of `from_raw_parts()` never returns `Err`.
+        <Self as SparseMerkleTree<SMT_DEPTH>>::from_raw_parts(inner_nodes, leaves, root).unwrap()
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -284,6 +301,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 root(&self) -> RpoDigest {
         self.root
     }

src/merkle/smt/mod.rs

@@ -292,6 +292,16 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     // 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<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,23 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
         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<NodeIndex, InnerNode>,
+        leaves: BTreeMap<u64, Word>,
+        root: RpoDigest,
+    ) -> Self {
+        // Our particular implementation of `from_raw_parts()` never returns `Err`.
+        <Self as SparseMerkleTree<DEPTH>>::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<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 root(&self) -> RpoDigest {
         self.root
     }