From 76eb7f6d0186d4c75a47411a2bfd7be407599eda Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Tue, 5 Nov 2024 16:17:41 -0700
Subject: [PATCH] smt: add from_raw_parts() to trait interface

---
 src/merkle/smt/full/mod.rs   | 21 +++++++++++++++++++++
 src/merkle/smt/mod.rs        |  8 ++++++++
 src/merkle/smt/simple/mod.rs | 21 +++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 0b170d6..34a677c 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/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> {
diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index 7ef302e..22b6d53 100644
--- a/src/merkle/smt/mod.rs
+++ b/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;
 
diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs
index 0a30b04..62ce753 100644
--- a/src/merkle/smt/simple/mod.rs
+++ b/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> {