diff --git a/src/merkle/empty_roots.rs b/src/merkle/empty_roots.rs index 17cd781..724d8ec 100644 --- a/src/merkle/empty_roots.rs +++ b/src/merkle/empty_roots.rs @@ -10,12 +10,19 @@ pub struct EmptySubtreeRoots; 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 @@ fn all_depths_opens_to_zero() { .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]); + } +} diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index a951122..1e3241c 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -19,7 +19,6 @@ pub struct SimpleSmt { root: RpoDigest, leaves: BTreeMap, branches: BTreeMap, - empty_hashes: Vec, } impl SimpleSmt { @@ -52,13 +51,11 @@ impl SimpleSmt { 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 @@ impl SimpleSmt { } 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 @@ impl SimpleSmt { 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 } }) }