M src/merkle/empty_roots.rs → src/merkle/empty_roots.rs

@@ -10,12 +10,19 @@
 impl EmptySubtreeRoots {
     /// Returns a static slice with roots of empty subtrees of a Merkle tree starting at the
     /// specified depth.
-    pub const fn empty_hashes(depth: u8) -> &'static [RpoDigest] {
-        let ptr = &EMPTY_SUBTREES[255 - depth as usize] as *const RpoDigest;
+    pub const fn empty_hashes(tree_depth: u8) -> &'static [RpoDigest] {
+        let ptr = &EMPTY_SUBTREES[255 - tree_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) }
+        unsafe { slice::from_raw_parts(ptr, tree_depth as usize + 1) }
+    }
+
+    /// Returns the node's digest for a sub-tree with all its leaves set to the empty word.
+    pub const fn entry(tree_depth: u8, node_depth: u8) -> &'static RpoDigest {
+        assert!(node_depth <= tree_depth);
+        let pos = 255 - tree_depth + node_depth;
+        &EMPTY_SUBTREES[pos as usize]
     }
 }
 
@@ -1583,3 +1590,16 @@ })
             .for_each(|(x, computed)| assert_eq!(x, computed));
     }
 }
+
+#[test]
+fn test_entry() {
+    // check the leaf is always the empty work
+    for depth in 0..255 {
+        assert_eq!(EmptySubtreeRoots::entry(depth, depth), &RpoDigest::new(EMPTY_WORD));
+    }
+
+    // check the root matches the first element of empty_hashes
+    for depth in 0..255 {
+        assert_eq!(EmptySubtreeRoots::entry(depth, 0), &EmptySubtreeRoots::empty_hashes(depth)[0]);
+    }
+}

M src/merkle/simple_smt/mod.rs → src/merkle/simple_smt/mod.rs

@@ -19,7 +19,6 @@ depth: u8,
     root: RpoDigest,
     leaves: BTreeMap<u64, Word>,
     branches: BTreeMap<NodeIndex, BranchNode>,
-    empty_hashes: Vec<RpoDigest>,
 }
 
 impl SimpleSmt {
@@ -52,13 +51,11 @@ } else if Self::MAX_DEPTH < depth {
             return Err(MerkleError::DepthTooBig(depth as u64));
         }
 
-        let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec();
-        let root = empty_hashes[0];
+        let root = *EmptySubtreeRoots::entry(depth, 0);
 
         Ok(Self {
             root,
             depth,
-            empty_hashes,
             leaves: BTreeMap::new(),
             branches: BTreeMap::new(),
         })
@@ -133,10 +130,12 @@ Err(MerkleError::DepthTooBig(index.depth() as u64))
         } else if index.depth() == self.depth() {
             // the lookup in empty_hashes could fail only if empty_hashes were not built correctly
             // by the constructor as we check the depth of the lookup above.
-            Ok(RpoDigest::from(
-                self.get_leaf_node(index.value())
-                    .unwrap_or_else(|| *self.empty_hashes[index.depth() as usize]),
-            ))
+            let leaf_pos = index.value();
+            let leaf = match self.get_leaf_node(leaf_pos) {
+                Some(word) => word.into(),
+                None => *EmptySubtreeRoots::entry(self.depth, index.depth()),
+            };
+            Ok(leaf)
         } else {
             Ok(self.get_branch_node(&index).parent())
         }
@@ -248,8 +247,8 @@ }
 
     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 }
+            let node = EmptySubtreeRoots::entry(self.depth, index.depth() + 1);
+            BranchNode { left: *node, right: *node }
         })
     }