Move leaf node creation into separate fn

eth2/utils/cached_tree_hash/src/btree_overlay.rs

@@ -37,7 +37,6 @@ pub struct BTreeOverlay {
     offset: usize,
     pub depth: usize,
     lengths: Vec<usize>,
-    leaf_nodes: Vec<LeafNode>,
 }
 
 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();
+        Self {
-        let num_internal_nodes = num_leaf_nodes - 1;
+            offset,
+            depth: schema.depth,
+            lengths: schema.lengths,
+        }
+    }
 
-        let mut running_offset = offset + num_internal_nodes;
+    pub fn get_leaf_nodes(&self, n: usize) -> Vec<LeafNode> {
-        let leaf_nodes: Vec<LeafNode> = schema
+        let mut running_offset = self.offset + self.num_internal_nodes();
+
+        let mut leaf_nodes: Vec<LeafNode> = self
             .lengths
             .iter()
             .map(|length| {
@@ -63,12 +68,10 @@ impl BTreeOverlay {
             })
             .collect();
 
-        Self {
+        leaf_nodes.resize(self.num_leaf_nodes(), LeafNode::Padding);
-            offset,
+        leaf_nodes.resize(n, LeafNode::DoesNotExist);
-            depth: schema.depth,
+
-            lengths: schema.lengths,
+        leaf_nodes
-            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<LeafNode, Error> {
-        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!(leaves[0], LeafNode::Exists(3..4));
-        assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5)));
+        assert_eq!(leaves[1], LeafNode::Exists(4..5));
-        assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6)));
+        assert_eq!(leaves[2], LeafNode::Exists(5..6));
-        assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Exists(6..7)));
+        assert_eq!(leaves[3], LeafNode::Exists(6..7));
-        assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist));
+        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!(leaves[0], LeafNode::Exists(3..4));
-        assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5)));
+        assert_eq!(leaves[1], LeafNode::Exists(4..5));
-        assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6)));
+        assert_eq!(leaves[2], LeafNode::Exists(5..6));
-        assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Padding));
+        assert_eq!(leaves[3], LeafNode::Padding);
-        assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist));
+        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!(leaves[0], LeafNode::Padding);
-        assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist));
+        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)));
+        let leaves = tree.get_leaf_nodes(2);
-        assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist));
+        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)));
+        let leaves = tree.get_leaf_nodes(2);
-        assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist));
+        assert_eq!(leaves[0], LeafNode::Exists(10..13));
+        assert_eq!(leaves[1], LeafNode::DoesNotExist);
     }
 
     #[test]

eth2/utils/cached_tree_hash/src/impls/vec.rs

@@ -147,11 +147,18 @@ pub fn update_tree_hash_cache<T: CachedTreeHash<T>>(
             }
         }
         TreeHashType::Container | TreeHashType::List | TreeHashType::Vector => {
-            for i in 0..std::cmp::max(new_overlay.num_leaf_nodes(), old_overlay.num_leaf_nodes()) {
+            let longest_len =
-                match (
+                std::cmp::max(new_overlay.num_leaf_nodes(), old_overlay.num_leaf_nodes());
-                    old_overlay.get_leaf_node(i + old_overlay.num_internal_nodes())?,
+
-                    new_overlay.get_leaf_node(i + new_overlay.num_internal_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.