WIP: slightly refactor recompute_nodes_from_index_to_root()

This commit is contained in:
Qyriad 2024-08-16 19:22:15 -06:00
parent d5250cd202
commit 2fac7f11d9

View file

@ -116,27 +116,24 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
mut index: NodeIndex, mut index: NodeIndex,
node_hash_at_index: RpoDigest, node_hash_at_index: RpoDigest,
) { ) {
let mut node_hash = node_hash_at_index; let mut new_child_hash = node_hash_at_index;
for node_depth in (0..index.depth()).rev() { for node_depth in (0..index.depth()).rev() {
let is_right = index.is_value_odd(); let is_right = index.is_value_odd();
index.move_up(); index.move_up();
let InnerNode { left, right } = self.get_inner_node(index); let old_node = self.get_inner_node(index);
let (left, right) = if is_right { let new_node = old_node.create_update(is_right, new_child_hash);
(left, node_hash) new_child_hash = new_node.hash();
} else {
(node_hash, right)
};
node_hash = Rpo256::merge(&[left, right]);
if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) { let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
if new_child_hash == equivalent_empty_hash {
// If a subtree is empty, when can remove the inner node, since it's equal to the // If a subtree is empty, when can remove the inner node, since it's equal to the
// default value // default value
self.remove_inner_node(index) self.remove_inner_node(index)
} else { } else {
self.insert_inner_node(index, InnerNode { left, right }); self.insert_inner_node(index, new_node);
} }
} }
self.set_root(node_hash); self.set_root(new_child_hash);
} }
// REQUIRED METHODS // REQUIRED METHODS
@ -195,6 +192,17 @@ pub(crate) struct InnerNode {
} }
impl InnerNode { impl InnerNode {
/// Creates a new node for when a child node has been changed.
pub fn create_update(self, update_right: bool, new_hash: RpoDigest) -> Self {
let (left, right) = if update_right {
(self.left, new_hash)
} else {
(new_hash, self.right)
};
Self { left, right }
}
pub fn hash(&self) -> RpoDigest { pub fn hash(&self) -> RpoDigest {
Rpo256::merge(&[self.left, self.right]) Rpo256::merge(&[self.left, self.right])
} }