Merge pull request #87 from 0xPolygonMiden/vlopes11-36-simple-smt-prepare

feat: refactor simple smt to use empty subtree constants

benches/smt.rs

@@ -18,7 +18,7 @@ fn smt_rpo(c: &mut Criterion) {
                     (i, word)
                 })
                 .collect();
-            let tree = SimpleSmt::new(entries, depth).unwrap();
+            let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap();
             trees.push(tree);
         }
     }

src/hash/rpo/digest.rs

@@ -11,7 +11,7 @@ use core::{cmp::Ordering, ops::Deref};
 pub struct RpoDigest([Felt; DIGEST_SIZE]);
 
 impl RpoDigest {
-    pub fn new(value: [Felt; DIGEST_SIZE]) -> Self {
+    pub const fn new(value: [Felt; DIGEST_SIZE]) -> Self {
         Self(value)
     }
 

src/merkle/empty_roots.rs

@@ -0,0 +1,454 @@
+use super::{Felt, RpoDigest, WORD_SIZE, ZERO};
+use core::slice;
+
+// EMPTY NODES SUBTREES
+// ================================================================================================
+
+/// Contains precomputed roots of empty subtrees in a Merkle rtee of depth 64.
+pub struct EmptySubtreeRoots;
+
+impl EmptySubtreeRoots {
+    /// Returns a static slice with roots of empty subtrees of a Merkle tree starting at the
+    /// specified depth.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the provided `depth` is greater than `64`.
+    pub const fn empty_hashes(depth: u8) -> &'static [RpoDigest] {
+        assert!(depth < 65);
+        let ptr = &EMPTY_SUBTREES_64[64 - depth as usize] as *const RpoDigest;
+        // Safety: this is a static/constant array, so it will never be outlived. If we attempt to
+        // use regular slices, this wouldn't be a `const` function, meaning we won't be able to use
+        // the returned value for static/constant definitions.
+        unsafe { slice::from_raw_parts(ptr, depth as usize + 1) }
+    }
+}
+
+const EMPTY_SUBTREES_64: [RpoDigest; 65] = [
+    RpoDigest::new([
+        Felt::new(15321474589252129342),
+        Felt::new(17373224439259377994),
+        Felt::new(15071539326562317628),
+        Felt::new(3312677166725950353),
+    ]),
+    RpoDigest::new([
+        Felt::new(12146678323567200178),
+        Felt::new(14288630174929498478),
+        Felt::new(13374892366980833045),
+        Felt::new(11840636859983936891),
+    ]),
+    RpoDigest::new([
+        Felt::new(15220380953028059006),
+        Felt::new(2981707349961006045),
+        Felt::new(7409523958661360004),
+        Felt::new(2816116826688969892),
+    ]),
+    RpoDigest::new([
+        Felt::new(7829641133220670678),
+        Felt::new(6170216088031698405),
+        Felt::new(11814483661801576435),
+        Felt::new(1762887097744793975),
+    ]),
+    RpoDigest::new([
+        Felt::new(1299421782687082884),
+        Felt::new(9938699043036414489),
+        Felt::new(10193025806762503939),
+        Felt::new(12073246492422971113),
+    ]),
+    RpoDigest::new([
+        Felt::new(3774016405860870757),
+        Felt::new(2584714598467121158),
+        Felt::new(7418645462301488344),
+        Felt::new(1016804897028793820),
+    ]),
+    RpoDigest::new([
+        Felt::new(13238072489118494737),
+        Felt::new(6917129315345826393),
+        Felt::new(13736362398490889690),
+        Felt::new(4929049375601714136),
+    ]),
+    RpoDigest::new([
+        Felt::new(2433738165854950976),
+        Felt::new(6710644905925382197),
+        Felt::new(10571480102433401045),
+        Felt::new(16853295309134271298),
+    ]),
+    RpoDigest::new([
+        Felt::new(3162775558610426184),
+        Felt::new(11944004899624546116),
+        Felt::new(55767976185223284),
+        Felt::new(5892480272697245897),
+    ]),
+    RpoDigest::new([
+        Felt::new(12582634330812132159),
+        Felt::new(6886254574119140332),
+        Felt::new(4407453795368410417),
+        Felt::new(6959805977831121004),
+    ]),
+    RpoDigest::new([
+        Felt::new(16001070406220863863),
+        Felt::new(4426773743735082930),
+        Felt::new(6860108527212616559),
+        Felt::new(3994703491288516722),
+    ]),
+    RpoDigest::new([
+        Felt::new(9755907048710665826),
+        Felt::new(13697078808748604851),
+        Felt::new(17210321635283113095),
+        Felt::new(1203394006092675979),
+    ]),
+    RpoDigest::new([
+        Felt::new(3332855817731547893),
+        Felt::new(1068928372599561798),
+        Felt::new(17119375903210334455),
+        Felt::new(8148601736624954416),
+    ]),
+    RpoDigest::new([
+        Felt::new(17265634841675424144),
+        Felt::new(18322832739735580203),
+        Felt::new(17896992777163902308),
+        Felt::new(6189383326950297131),
+    ]),
+    RpoDigest::new([
+        Felt::new(9329637674239983584),
+        Felt::new(2512861030579248721),
+        Felt::new(10833150484884266896),
+        Felt::new(7470498642428983444),
+    ]),
+    RpoDigest::new([
+        Felt::new(3611140194800558886),
+        Felt::new(17185933650781587767),
+        Felt::new(7835232399818923215),
+        Felt::new(7974155618002781326),
+    ]),
+    RpoDigest::new([
+        Felt::new(17483286922353768131),
+        Felt::new(353378057542380712),
+        Felt::new(1935183237414585408),
+        Felt::new(4820339620987989650),
+    ]),
+    RpoDigest::new([
+        Felt::new(16172462385444809646),
+        Felt::new(3268597753131435459),
+        Felt::new(3481491333654579291),
+        Felt::new(16487779176137683725),
+    ]),
+    RpoDigest::new([
+        Felt::new(16595012576192613315),
+        Felt::new(16028552537812484518),
+        Felt::new(13016887826405546773),
+        Felt::new(14649690775021494057),
+    ]),
+    RpoDigest::new([
+        Felt::new(11300236651178143890),
+        Felt::new(15307634289168527196),
+        Felt::new(2834866419963148279),
+        Felt::new(7512874625395280090),
+    ]),
+    RpoDigest::new([
+        Felt::new(1148273481270068529),
+        Felt::new(7411276436636897120),
+        Felt::new(14325955409748352141),
+        Felt::new(15577038614919538356),
+    ]),
+    RpoDigest::new([
+        Felt::new(13911627859049081064),
+        Felt::new(13298542751859672529),
+        Felt::new(18341014824837028242),
+        Felt::new(5587966507704160144),
+    ]),
+    RpoDigest::new([
+        Felt::new(10957185917743597702),
+        Felt::new(15815185767119166433),
+        Felt::new(17883994521792846784),
+        Felt::new(15958104556930886663),
+    ]),
+    RpoDigest::new([
+        Felt::new(13148367538964199489),
+        Felt::new(7372139436485928380),
+        Felt::new(13408383191801051600),
+        Felt::new(2114382634401123096),
+    ]),
+    RpoDigest::new([
+        Felt::new(14448157482521530067),
+        Felt::new(17865161921504959156),
+        Felt::new(10319385198642448897),
+        Felt::new(364163501511998552),
+    ]),
+    RpoDigest::new([
+        Felt::new(9722640569118951143),
+        Felt::new(16371655672847089887),
+        Felt::new(12379452272155069993),
+        Felt::new(11605969747977185617),
+    ]),
+    RpoDigest::new([
+        Felt::new(2782512273606877924),
+        Felt::new(3656296563981095117),
+        Felt::new(5947388149010135441),
+        Felt::new(1678144343036748885),
+    ]),
+    RpoDigest::new([
+        Felt::new(10347491038074052866),
+        Felt::new(11061756013655443653),
+        Felt::new(8901792852813329415),
+        Felt::new(10002477867799577447),
+    ]),
+    RpoDigest::new([
+        Felt::new(16688151588649906570),
+        Felt::new(12937054427339650762),
+        Felt::new(2125115528195796454),
+        Felt::new(4796610823085621719),
+    ]),
+    RpoDigest::new([
+        Felt::new(3032620037225059051),
+        Felt::new(13522881885116127385),
+        Felt::new(6010511038055304264),
+        Felt::new(8199256447383686121),
+    ]),
+    RpoDigest::new([
+        Felt::new(11250302734399433639),
+        Felt::new(4970037623163209776),
+        Felt::new(15776613712371118341),
+        Felt::new(5554382612311754837),
+    ]),
+    RpoDigest::new([
+        Felt::new(5116523511540088640),
+        Felt::new(12381059245485642368),
+        Felt::new(2176361879916914688),
+        Felt::new(11209293198464735683),
+    ]),
+    RpoDigest::new([
+        Felt::new(11677748883385181208),
+        Felt::new(15891398395707500576),
+        Felt::new(3790704659934033620),
+        Felt::new(2126099371106695189),
+    ]),
+    RpoDigest::new([
+        Felt::new(13948603355603496603),
+        Felt::new(15902438544472945077),
+        Felt::new(1969361494026622497),
+        Felt::new(17326911676634210553),
+    ]),
+    RpoDigest::new([
+        Felt::new(16081431322775411514),
+        Felt::new(13201312030265587002),
+        Felt::new(18283434127959076535),
+        Felt::new(9889802180847551599),
+    ]),
+    RpoDigest::new([
+        Felt::new(8490051641633132830),
+        Felt::new(11985660456681176415),
+        Felt::new(12193381039977027251),
+        Felt::new(17563185381678568385),
+    ]),
+    RpoDigest::new([
+        Felt::new(3870617340693651786),
+        Felt::new(2748490321246408799),
+        Felt::new(8501743976565218963),
+        Felt::new(1660720190266083389),
+    ]),
+    RpoDigest::new([
+        Felt::new(2121119282758520982),
+        Felt::new(9042267662074029772),
+        Felt::new(15431993929052434204),
+        Felt::new(10659345458998811701),
+    ]),
+    RpoDigest::new([
+        Felt::new(15206763021853065070),
+        Felt::new(15268692497656424421),
+        Felt::new(13335448435922172445),
+        Felt::new(3421340628484408379),
+    ]),
+    RpoDigest::new([
+        Felt::new(5175159910654039438),
+        Felt::new(10258564296733764665),
+        Felt::new(235961379704359454),
+        Felt::new(18007006485615491006),
+    ]),
+    RpoDigest::new([
+        Felt::new(9455184082727641653),
+        Felt::new(6634498452861935579),
+        Felt::new(18189776179964984407),
+        Felt::new(3546641211720870472),
+    ]),
+    RpoDigest::new([
+        Felt::new(2566088177506289568),
+        Felt::new(7785941571143323572),
+        Felt::new(13948908169667863201),
+        Felt::new(8557252288425473395),
+    ]),
+    RpoDigest::new([
+        Felt::new(8801845050152766755),
+        Felt::new(514652983374395586),
+        Felt::new(13975919271481418443),
+        Felt::new(17480955484347349170),
+    ]),
+    RpoDigest::new([
+        Felt::new(7078477424334594989),
+        Felt::new(9975053207879493059),
+        Felt::new(5220656123503260168),
+        Felt::new(13795787984352794188),
+    ]),
+    RpoDigest::new([
+        Felt::new(1478357986561897612),
+        Felt::new(3963701567400985039),
+        Felt::new(10269836564499521403),
+        Felt::new(11874873630603798755),
+    ]),
+    RpoDigest::new([
+        Felt::new(936391814816943993),
+        Felt::new(6085855616346025677),
+        Felt::new(5782721339195502211),
+        Felt::new(10409491632083436908),
+    ]),
+    RpoDigest::new([
+        Felt::new(11138475264090866271),
+        Felt::new(17799626597540451271),
+        Felt::new(17968790388406362807),
+        Felt::new(9539434947296310791),
+    ]),
+    RpoDigest::new([
+        Felt::new(13051724588530357940),
+        Felt::new(8058102530250142518),
+        Felt::new(1861782711432586670),
+        Felt::new(2928050228215055187),
+    ]),
+    RpoDigest::new([
+        Felt::new(10650694022550988030),
+        Felt::new(5634734408638476525),
+        Felt::new(9233115969432897632),
+        Felt::new(1437907447409278328),
+    ]),
+    RpoDigest::new([
+        Felt::new(9720135276484706819),
+        Felt::new(9350120041401976641),
+        Felt::new(1348777594376050933),
+        Felt::new(13138246165242825648),
+    ]),
+    RpoDigest::new([
+        Felt::new(10866643979409126085),
+        Felt::new(13790633638103642042),
+        Felt::new(6374461622011119670),
+        Felt::new(5702679962735491362),
+    ]),
+    RpoDigest::new([
+        Felt::new(5257277882444261955),
+        Felt::new(8511211402794551302),
+        Felt::new(3294838877645533839),
+        Felt::new(4084864647832858048),
+    ]),
+    RpoDigest::new([
+        Felt::new(7948776578097466250),
+        Felt::new(8630046431048474853),
+        Felt::new(11549811661672434609),
+        Felt::new(14329713552208961509),
+    ]),
+    RpoDigest::new([
+        Felt::new(734617692582477804),
+        Felt::new(11871516935077749937),
+        Felt::new(12085935336918533812),
+        Felt::new(11028098016323141988),
+    ]),
+    RpoDigest::new([
+        Felt::new(10937083382606895486),
+        Felt::new(12203867463821771187),
+        Felt::new(13369919265612777227),
+        Felt::new(2521482611471096233),
+    ]),
+    RpoDigest::new([
+        Felt::new(1242037330294600071),
+        Felt::new(8643213198640797337),
+        Felt::new(14112360612081236212),
+        Felt::new(11296904697431650998),
+    ]),
+    RpoDigest::new([
+        Felt::new(11958494925108187724),
+        Felt::new(6059642826232274823),
+        Felt::new(1563918267677757605),
+        Felt::new(266509853282035592),
+    ]),
+    RpoDigest::new([
+        Felt::new(17288335252189973373),
+        Felt::new(3243363076395469373),
+        Felt::new(8880515798614590986),
+        Felt::new(10260780639137628077),
+    ]),
+    RpoDigest::new([
+        Felt::new(1839714959437284152),
+        Felt::new(12088193186987715006),
+        Felt::new(10200898335013164008),
+        Felt::new(12768529781145127245),
+    ]),
+    RpoDigest::new([
+        Felt::new(1537615626967151439),
+        Felt::new(11731506816677487155),
+        Felt::new(4748463589169553420),
+        Felt::new(17495851576537541106),
+    ]),
+    RpoDigest::new([
+        Felt::new(957733314860117562),
+        Felt::new(15623410588944187169),
+        Felt::new(4321611031548662227),
+        Felt::new(12856104259650439278),
+    ]),
+    RpoDigest::new([
+        Felt::new(14827447693720375746),
+        Felt::new(17296925942589213350),
+        Felt::new(13524332314559504765),
+        Felt::new(15663886706087995199),
+    ]),
+    RpoDigest::new([
+        Felt::new(18185978518863914335),
+        Felt::new(936586966360019113),
+        Felt::new(497299419609993926),
+        Felt::new(1977881506773614749),
+    ]),
+    RpoDigest::new([
+        Felt::new(8635338869442206704),
+        Felt::new(11671305615285950885),
+        Felt::new(15253023094703789604),
+        Felt::new(7398108415970215319),
+    ]),
+    RpoDigest::new([ZERO; WORD_SIZE]),
+];
+
+#[test]
+fn all_depths_opens_to_zero() {
+    use super::Rpo256;
+
+    for depth in 1..=64 {
+        // fetch the subtrees and reverse it so the path is leaf -> root
+        let mut subtree = EmptySubtreeRoots::empty_hashes(depth).to_vec();
+        subtree.reverse();
+
+        // the length of the subtrees set must be equal to depth + 1 as we also
+        // include the root
+        assert_eq!(depth as usize + 1, subtree.len());
+
+        // assert the opening is zero
+        let initial = RpoDigest::new([ZERO; WORD_SIZE]);
+        assert_eq!(initial, subtree.remove(0));
+
+        // compute every node of the path manually and compare with the output
+        subtree
+            .into_iter()
+            .scan(initial, |state, x| {
+                *state = Rpo256::merge(&[*state; 2]);
+                Some((x, *state))
+            })
+            .for_each(|(x, computed)| assert_eq!(x, computed));
+    }
+}
+
+#[test]
+fn arbitrary_inputs_will_generate_sound_slices() {
+    let min = &EMPTY_SUBTREES_64[0] as *const RpoDigest;
+    let max = unsafe { min.add(64) };
+    for depth in 0..=64 {
+        let subtree = EmptySubtreeRoots::empty_hashes(depth);
+        let first = &subtree[0] as *const RpoDigest;
+        let last = &subtree[depth as usize] as *const RpoDigest;
+        assert!(min <= first && first <= max);
+        assert!(min <= last && last <= max);
+    }
+}

src/merkle/mod.rs

@@ -1,12 +1,16 @@
 use super::{
     hash::rpo::{Rpo256, RpoDigest},
     utils::collections::{vec, BTreeMap, Vec},
-    Felt, StarkField, Word, ZERO,
+    Felt, StarkField, Word, WORD_SIZE, ZERO,
 };
 use core::fmt;
 
 // REEXPORTS
 // ================================================================================================
+
+mod empty_roots;
+pub use empty_roots::EmptySubtreeRoots;
+
 mod index;
 pub use index::NodeIndex;
 

src/merkle/simple_smt/mod.rs

@@ -1,4 +1,6 @@
-use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
+use super::{
+    BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word,
+};
 
 #[cfg(test)]
 mod tests;
@@ -29,38 +31,55 @@ impl SimpleSmt {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    /// Creates a new simple SMT.
-    ///
-    /// The provided entries will be tuples of the leaves and their corresponding keys.
+    /// Creates a new simple SMT with the provided depth.
+    pub fn new(depth: u8) -> Result<Self, MerkleError> {
+        // validate the range of the depth.
+        if depth < Self::MIN_DEPTH {
+            return Err(MerkleError::DepthTooSmall(depth));
+        } else if Self::MAX_DEPTH < depth {
+            return Err(MerkleError::DepthTooBig(depth as u64));
+        }
+
+        let (store, root) = Store::new(depth);
+        Ok(Self { root, depth, store })
+    }
+
+    /// Appends the provided entries as leaves of the tree.
     ///
     /// # Errors
     ///
     /// The function will fail if the provided entries count exceed the maximum tree capacity, that
     /// is `2^{depth}`.
-    pub fn new<R, I>(entries: R, depth: u8) -> Result<Self, MerkleError>
+    pub fn with_leaves<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
     where
         R: IntoIterator<IntoIter = I>,
         I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
     {
+        // check if the leaves count will fit the depth setup
         let mut entries = entries.into_iter();
-
-        // validate the range of the depth.
-        let max = 1 << depth;
-        if depth < Self::MIN_DEPTH {
-            return Err(MerkleError::DepthTooSmall(depth));
-        } else if Self::MAX_DEPTH < depth {
-            return Err(MerkleError::DepthTooBig(depth as u64));
-        } else if entries.len() > max {
+        let max = 1 << self.depth;
+        if entries.len() > max {
             return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
         }
 
-        let (store, root) = Store::new(depth);
-        let mut tree = Self { root, depth, store };
-        entries.try_for_each(|(key, leaf)| tree.insert_leaf(key, leaf))?;
+        // append leaves and return
+        entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?;
+        Ok(self)
+    }
 
-        Ok(tree)
+    /// Replaces the internal empty digests used when a given depth doesn't contain a node.
+    pub fn with_empty_subtrees<I>(mut self, hashes: I) -> Self
+    where
+        I: IntoIterator<Item = RpoDigest>,
+    {
+        self.store
+            .replace_empty_subtrees(hashes.into_iter().collect());
+        self
     }
 
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
     /// Returns the root of this Merkle tree.
     pub const fn root(&self) -> Word {
         self.root
@@ -71,6 +90,9 @@ impl SimpleSmt {
         self.depth
     }
 
+    // PROVIDERS
+    // --------------------------------------------------------------------------------------------
+
     /// Returns the set count of the keys of the leaves.
     pub fn leaves_count(&self) -> usize {
         self.store.leaves_count()
@@ -81,16 +103,24 @@ impl SimpleSmt {
     /// # Errors
     /// Returns an error if:
     /// * The specified depth is greater than the depth of the tree.
-    /// * The specified key does not exist
     pub fn get_node(&self, index: &NodeIndex) -> Result<Word, MerkleError> {
         if index.is_root() {
             Err(MerkleError::DepthTooSmall(index.depth()))
         } else if index.depth() > self.depth() {
             Err(MerkleError::DepthTooBig(index.depth() as u64))
         } else if index.depth() == self.depth() {
-            self.store.get_leaf_node(index.value())
+            self.store
+                .get_leaf_node(index.value())
+                .or_else(|| {
+                    self.store
+                        .empty_hashes
+                        .get(index.depth() as usize)
+                        .copied()
+                        .map(Word::from)
+                })
+                .ok_or(MerkleError::InvalidIndex(*index))
         } else {
-            let branch_node = self.store.get_branch_node(index)?;
+            let branch_node = self.store.get_branch_node(index);
             Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
         }
     }
@@ -100,23 +130,19 @@ impl SimpleSmt {
     ///
     /// # Errors
     /// Returns an error if:
-    /// * The specified key does not exist as a branch or leaf node
     /// * The specified depth is greater than the depth of the tree.
     pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
         if index.is_root() {
             return Err(MerkleError::DepthTooSmall(index.depth()));
         } else if index.depth() > self.depth() {
             return Err(MerkleError::DepthTooBig(index.depth() as u64));
-        } else if index.depth() == self.depth() && !self.store.check_leaf_node_exists(index.value())
-        {
-            return Err(MerkleError::InvalidIndex(index.with_depth(self.depth())));
         }
 
         let mut path = Vec::with_capacity(index.depth() as usize);
         for _ in 0..index.depth() {
             let is_right = index.is_value_odd();
             index.move_up();
-            let BranchNode { left, right } = self.store.get_branch_node(&index)?;
+            let BranchNode { left, right } = self.store.get_branch_node(&index);
             let value = if is_right { left } else { right };
             path.push(*value);
         }
@@ -133,6 +159,9 @@ impl SimpleSmt {
         self.get_path(NodeIndex::new(self.depth(), key))
     }
 
+    // STATE MUTATORS
+    // --------------------------------------------------------------------------------------------
+
     /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
     ///
     /// # Errors
@@ -156,10 +185,7 @@ impl SimpleSmt {
         for _ in 0..index.depth() {
             let is_right = index.is_value_odd();
             index.move_up();
-            let BranchNode { left, right } = self
-                .store
-                .get_branch_node(&index)
-                .unwrap_or_else(|_| self.store.get_empty_node(index.depth() as usize + 1));
+            let BranchNode { left, right } = self.store.get_branch_node(&index);
             let (left, right) = if is_right {
                 (left, value)
             } else {
@@ -200,16 +226,7 @@ impl Store {
         let leaves = BTreeMap::new();
 
         // Construct empty node digests for each layer of the tree
-        let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
-            .scan(Word::default().into(), |state, _| {
-                let value = *state;
-                *state = Rpo256::merge(&[value, value]);
-                Some(value)
-            })
-            .collect::<Vec<_>>()
-            .into_iter()
-            .rev()
-            .collect();
+        let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec();
 
         let root = empty_hashes[0].into();
         let store = Self {
@@ -222,34 +239,30 @@ impl Store {
         (store, root)
     }
 
-    fn get_empty_node(&self, depth: usize) -> BranchNode {
-        let digest = self.empty_hashes[depth];
-        BranchNode {
-            left: digest,
-            right: digest,
-        }
+    fn replace_empty_subtrees(&mut self, hashes: Vec<RpoDigest>) {
+        self.empty_hashes = hashes;
     }
 
     fn check_leaf_node_exists(&self, key: u64) -> bool {
         self.leaves.contains_key(&key)
     }
 
-    fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
-        self.leaves
-            .get(&key)
-            .cloned()
-            .ok_or(MerkleError::InvalidIndex(NodeIndex::new(self.depth, key)))
+    fn get_leaf_node(&self, key: u64) -> Option<Word> {
+        self.leaves.get(&key).copied()
     }
 
     fn insert_leaf_node(&mut self, key: u64, node: Word) {
         self.leaves.insert(key, node);
     }
 
-    fn get_branch_node(&self, index: &NodeIndex) -> Result<BranchNode, MerkleError> {
-        self.branches
-            .get(index)
-            .cloned()
-            .ok_or(MerkleError::InvalidIndex(*index))
+    fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
+        self.branches.get(index).cloned().unwrap_or_else(|| {
+            let node = self.empty_hashes[index.depth() as usize + 1];
+            BranchNode {
+                left: node,
+                right: node,
+            }
+        })
     }
 
     fn insert_branch_node(&mut self, index: NodeIndex, left: RpoDigest, right: RpoDigest) {

src/merkle/simple_smt/tests.rs

@@ -2,7 +2,6 @@ use super::{
     super::{int_to_node, MerkleTree, RpoDigest, SimpleSmt},
     NodeIndex, Rpo256, Vec, Word,
 };
-use core::iter;
 use proptest::prelude::*;
 use rand_utils::prng_array;
 
@@ -31,7 +30,7 @@ const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
 
 #[test]
 fn build_empty_tree() {
-    let smt = SimpleSmt::new(iter::empty(), 3).unwrap();
+    let smt = SimpleSmt::new(3).unwrap();
     let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
     assert_eq!(mt.root(), smt.root());
 }
@@ -39,7 +38,7 @@ fn build_empty_tree() {
 #[test]
 fn empty_digests_are_consistent() {
     let depth = 5;
-    let root = SimpleSmt::new(iter::empty(), depth).unwrap().root();
+    let root = SimpleSmt::new(depth).unwrap().root();
     let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| {
         let digest = Rpo256::merge(&state);
         [digest; 2]
@@ -50,7 +49,7 @@ fn empty_digests_are_consistent() {
 
 #[test]
 fn build_sparse_tree() {
-    let mut smt = SimpleSmt::new(iter::empty(), 3).unwrap();
+    let mut smt = SimpleSmt::new(3).unwrap();
     let mut values = ZERO_VALUES8.to_vec();
 
     // insert single value
@@ -82,7 +81,10 @@ fn build_sparse_tree() {
 
 #[test]
 fn build_full_tree() {
-    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
+    let tree = SimpleSmt::new(2)
+        .unwrap()
+        .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
+        .unwrap();
 
     let (root, node2, node3) = compute_internal_nodes();
     assert_eq!(root, tree.root());
@@ -92,7 +94,10 @@ fn build_full_tree() {
 
 #[test]
 fn get_values() {
-    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
+    let tree = SimpleSmt::new(2)
+        .unwrap()
+        .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
+        .unwrap();
 
     // check depth 2
     assert_eq!(VALUES4[0], tree.get_node(&NodeIndex::new(2, 0)).unwrap());
@@ -103,7 +108,10 @@ fn get_values() {
 
 #[test]
 fn get_path() {
-    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
+    let tree = SimpleSmt::new(2)
+        .unwrap()
+        .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
+        .unwrap();
 
     let (_, node2, node3) = compute_internal_nodes();
 
@@ -132,18 +140,20 @@ fn get_path() {
 
 #[test]
 fn update_leaf() {
-    let mut tree = SimpleSmt::new(KEYS8.into_iter().zip(VALUES8.into_iter()), 3).unwrap();
+    let mut tree = SimpleSmt::new(3)
+        .unwrap()
+        .with_leaves(KEYS8.into_iter().zip(VALUES8.into_iter()))
+        .unwrap();
 
     // update one value
     let key = 3;
     let new_node = int_to_node(9);
     let mut expected_values = VALUES8.to_vec();
     expected_values[key] = new_node;
-    let expected_tree = SimpleSmt::new(
-        KEYS8.into_iter().zip(expected_values.clone().into_iter()),
-        3,
-    )
-    .unwrap();
+    let expected_tree = SimpleSmt::new(3)
+        .unwrap()
+        .with_leaves(KEYS8.into_iter().zip(expected_values.clone().into_iter()))
+        .unwrap();
 
     tree.update_leaf(key as u64, new_node).unwrap();
     assert_eq!(expected_tree.root, tree.root);
@@ -152,8 +162,10 @@ fn update_leaf() {
     let key = 6;
     let new_node = int_to_node(10);
     expected_values[key] = new_node;
-    let expected_tree =
-        SimpleSmt::new(KEYS8.into_iter().zip(expected_values.into_iter()), 3).unwrap();
+    let expected_tree = SimpleSmt::new(3)
+        .unwrap()
+        .with_leaves(KEYS8.into_iter().zip(expected_values.into_iter()))
+        .unwrap();
 
     tree.update_leaf(key as u64, new_node).unwrap();
     assert_eq!(expected_tree.root, tree.root);
@@ -188,7 +200,7 @@ fn small_tree_opening_is_consistent() {
 
     let depth = 3;
     let entries = vec![(0, a), (1, b), (4, c), (7, d)];
-    let tree = SimpleSmt::new(entries, depth).unwrap();
+    let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap();
 
     assert_eq!(tree.root(), Word::from(k));
 
@@ -219,7 +231,7 @@ proptest! {
         key in prop::num::u64::ANY,
         leaf in prop::num::u64::ANY,
     ) {
-        let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
+        let mut tree = SimpleSmt::new(depth).unwrap();
 
         let key = key % (1 << depth as u64);
         let leaf = int_to_node(leaf);
@@ -240,7 +252,7 @@ proptest! {
         count in 2u8..10u8,
         ref seed in any::<[u8; 32]>()
     ) {
-        let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
+        let mut tree = SimpleSmt::new(depth).unwrap();
         let mut seed = *seed;
         let leaves = (1 << depth) - 1;