diff --git a/eth2/utils/cached_tree_hash/src/btree_overlay.rs b/eth2/utils/cached_tree_hash/src/btree_overlay.rs index 2ad2383c8..46b1f31b8 100644 --- a/eth2/utils/cached_tree_hash/src/btree_overlay.rs +++ b/eth2/utils/cached_tree_hash/src/btree_overlay.rs @@ -37,7 +37,6 @@ pub struct BTreeOverlay { offset: usize, pub depth: usize, lengths: Vec, - leaf_nodes: Vec, } impl BTreeOverlay { @@ -49,11 +48,17 @@ impl BTreeOverlay { } pub fn from_schema(schema: BTreeSchema, offset: usize) -> Self { - let num_leaf_nodes = schema.lengths.len().next_power_of_two(); - let num_internal_nodes = num_leaf_nodes - 1; + Self { + offset, + depth: schema.depth, + lengths: schema.lengths, + } + } - let mut running_offset = offset + num_internal_nodes; - let leaf_nodes: Vec = schema + pub fn get_leaf_nodes(&self, n: usize) -> Vec { + let mut running_offset = self.offset + self.num_internal_nodes(); + + let mut leaf_nodes: Vec = self .lengths .iter() .map(|length| { @@ -63,12 +68,10 @@ impl BTreeOverlay { }) .collect(); - Self { - offset, - depth: schema.depth, - lengths: schema.lengths, - leaf_nodes, - } + leaf_nodes.resize(self.num_leaf_nodes(), LeafNode::Padding); + leaf_nodes.resize(n, LeafNode::DoesNotExist); + + leaf_nodes } pub fn num_leaf_nodes(&self) -> usize { @@ -128,28 +131,6 @@ impl BTreeOverlay { self.offset + self.num_internal_nodes() } - /// Returns the chunk-range for a given leaf node. - /// - /// Returns `None` if: - /// - The specified node is internal. - /// - The specified node is padding. - /// - The specified node is OOB of the tree. - pub fn get_leaf_node(&self, i: usize) -> Result { - if i >= self.num_nodes() { - Ok(LeafNode::DoesNotExist) - } else if i >= self.num_nodes() - self.num_padding_leaves() { - Ok(LeafNode::Padding) - } else if (i == self.num_internal_nodes()) && (self.lengths.len() == 0) { - // If this is the first leaf node and the overlay contains zero items, return `None` as - // this node must be padding. - Ok(LeafNode::Padding) - } else { - let i = i - self.num_internal_nodes(); - - Ok(self.leaf_nodes[i].clone()) - } - } - pub fn child_chunks(&self, parent: usize) -> (usize, usize) { let children = children(parent); @@ -263,33 +244,38 @@ mod test { #[test] fn get_leaf_node() { let tree = get_tree_a(4); + let leaves = tree.get_leaf_nodes(5); - assert_eq!(tree.get_leaf_node(3), Ok(LeafNode::Exists(3..4))); - assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5))); - assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6))); - assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Exists(6..7))); - assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist)); + assert_eq!(leaves[0], LeafNode::Exists(3..4)); + assert_eq!(leaves[1], LeafNode::Exists(4..5)); + assert_eq!(leaves[2], LeafNode::Exists(5..6)); + assert_eq!(leaves[3], LeafNode::Exists(6..7)); + assert_eq!(leaves[4], LeafNode::DoesNotExist); let tree = get_tree_a(3); + let leaves = tree.get_leaf_nodes(5); - assert_eq!(tree.get_leaf_node(3), Ok(LeafNode::Exists(3..4))); - assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5))); - assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6))); - assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Padding)); - assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist)); + assert_eq!(leaves[0], LeafNode::Exists(3..4)); + assert_eq!(leaves[1], LeafNode::Exists(4..5)); + assert_eq!(leaves[2], LeafNode::Exists(5..6)); + assert_eq!(leaves[3], LeafNode::Padding); + assert_eq!(leaves[4], LeafNode::DoesNotExist); let tree = get_tree_a(0); + let leaves = tree.get_leaf_nodes(2); - assert_eq!(tree.get_leaf_node(0), Ok(LeafNode::Padding)); - assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist)); + assert_eq!(leaves[0], LeafNode::Padding); + assert_eq!(leaves[1], LeafNode::DoesNotExist); let tree = BTreeSchema::from_lengths(0, vec![3]).into_overlay(0); - assert_eq!(tree.get_leaf_node(0), Ok(LeafNode::Exists(0..3))); - assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist)); + let leaves = tree.get_leaf_nodes(2); + assert_eq!(leaves[0], LeafNode::Exists(0..3)); + assert_eq!(leaves[1], LeafNode::DoesNotExist); let tree = BTreeSchema::from_lengths(0, vec![3]).into_overlay(10); - assert_eq!(tree.get_leaf_node(0), Ok(LeafNode::Exists(10..13))); - assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist)); + let leaves = tree.get_leaf_nodes(2); + assert_eq!(leaves[0], LeafNode::Exists(10..13)); + assert_eq!(leaves[1], LeafNode::DoesNotExist); } #[test] diff --git a/eth2/utils/cached_tree_hash/src/impls/vec.rs b/eth2/utils/cached_tree_hash/src/impls/vec.rs index 046a1f97e..67a48ac0b 100644 --- a/eth2/utils/cached_tree_hash/src/impls/vec.rs +++ b/eth2/utils/cached_tree_hash/src/impls/vec.rs @@ -147,11 +147,18 @@ pub fn update_tree_hash_cache>( } } TreeHashType::Container | TreeHashType::List | TreeHashType::Vector => { - for i in 0..std::cmp::max(new_overlay.num_leaf_nodes(), old_overlay.num_leaf_nodes()) { - match ( - old_overlay.get_leaf_node(i + old_overlay.num_internal_nodes())?, - new_overlay.get_leaf_node(i + new_overlay.num_internal_nodes())?, - ) { + let longest_len = + std::cmp::max(new_overlay.num_leaf_nodes(), old_overlay.num_leaf_nodes()); + + let old_leaf_nodes = old_overlay.get_leaf_nodes(longest_len); + let new_leaf_nodes = if old_overlay == new_overlay { + old_leaf_nodes.clone() + } else { + new_overlay.get_leaf_nodes(longest_len) + }; + + for i in 0..longest_len { + match (&old_leaf_nodes[i], &new_leaf_nodes[i]) { // The item existed in the previous list and exists in the current list. // // Update the item.