Improve cached hashing performance

eth2/utils/cached_tree_hash/src/btree_overlay.rs

@@ -12,11 +12,7 @@ impl BTreeSchema {
     }
 
     pub fn into_overlay(self, offset: usize) -> BTreeOverlay {
-        BTreeOverlay {
+        BTreeOverlay::from_schema(self, offset)
-            offset,
-            depth: self.depth,
-            lengths: self.lengths,
-        }
     }
 }
 
@@ -41,6 +37,7 @@ pub struct BTreeOverlay {
     offset: usize,
     pub depth: usize,
     lengths: Vec<usize>,
+    leaf_nodes: Vec<LeafNode>,
 }
 
 impl BTreeOverlay {
@@ -48,8 +45,30 @@ impl BTreeOverlay {
     where
         T: CachedTreeHash<T>,
     {
-        item.tree_hash_cache_schema(depth)
+        Self::from_schema(item.tree_hash_cache_schema(depth), initial_offset)
-            .into_overlay(initial_offset)
+    }
+
+    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;
+
+        let mut running_offset = offset + num_internal_nodes;
+        let leaf_nodes: Vec<LeafNode> = schema
+            .lengths
+            .iter()
+            .map(|length| {
+                let range = running_offset..running_offset + length;
+                running_offset += length;
+                LeafNode::Exists(range)
+            })
+            .collect();
+
+        Self {
+            offset,
+            depth: schema.depth,
+            lengths: schema.lengths,
+            leaf_nodes,
+        }
     }
 
     pub fn num_leaf_nodes(&self) -> usize {
@@ -127,12 +146,7 @@ impl BTreeOverlay {
         } else {
             let i = i - self.num_internal_nodes();
 
-            let first_node = self.offset
+            Ok(self.leaf_nodes[i].clone())
-                + self.num_internal_nodes()
-                + self.lengths.iter().take(i).sum::<usize>();
-            let last_node = first_node + self.lengths[i];
-
-            Ok(LeafNode::Exists(first_node..last_node))
         }
     }
 
@@ -272,6 +286,10 @@ mod test {
         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 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));
     }
 
     #[test]

eth2/utils/cached_tree_hash/src/impls.rs

@@ -86,19 +86,22 @@ impl CachedTreeHash<[u8; 4]> for [u8; 4] {
 impl CachedTreeHash<H256> for H256 {
     fn new_tree_hash_cache(&self, _depth: usize) -> Result<TreeHashCache, Error> {
         Ok(TreeHashCache::from_bytes(
-            merkleize(self.as_bytes().to_vec()),
+            self.as_bytes().to_vec(),
             false,
             None,
         )?)
     }
 
+    fn num_tree_hash_cache_chunks(&self) -> usize {
+        1
+    }
+
     fn tree_hash_cache_schema(&self, depth: usize) -> BTreeSchema {
         BTreeSchema::from_lengths(depth, vec![1])
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {
-        let leaf = merkleize(self.as_bytes().to_vec());
+        cache.maybe_update_chunk(cache.chunk_index, self.as_bytes())?;
-        cache.maybe_update_chunk(cache.chunk_index, &leaf)?;
 
         cache.chunk_index += 1;
 

eth2/utils/cached_tree_hash/tests/tests.rs

@@ -1,7 +1,31 @@
 use cached_tree_hash::{merkleize::merkleize, *};
+use ethereum_types::H256 as Hash256;
 use int_to_bytes::int_to_bytes32;
 use tree_hash_derive::{CachedTreeHash, TreeHash};
 
+#[test]
+fn modifications() {
+    let n = 2048;
+
+    let vec: Vec<Hash256> = (0..n).map(|_| Hash256::random()).collect();
+
+    let mut cache = TreeHashCache::new(&vec, 0).unwrap();
+    cache.update(&vec).unwrap();
+
+    let modifications = cache.chunk_modified.iter().filter(|b| **b).count();
+
+    assert_eq!(modifications, 0);
+
+    let mut modified_vec = vec.clone();
+    modified_vec[n - 1] = Hash256::random();
+
+    cache.update(&modified_vec).unwrap();
+
+    let modifications = cache.chunk_modified.iter().filter(|b| **b).count();
+
+    assert_eq!(modifications, n.trailing_zeros() as usize + 2);
+}
+
 #[derive(Clone, Debug, TreeHash, CachedTreeHash)]
 pub struct NestedStruct {
     pub a: u64,
@@ -106,6 +130,24 @@ fn test_inner() {
     test_routine(original, modified);
 }
 
+#[test]
+fn test_vec_of_hash256() {
+    let n = 16;
+
+    let original: Vec<Hash256> = (0..n).map(|_| Hash256::random()).collect();
+
+    let modified: Vec<Vec<Hash256>> = vec![
+        original[..].to_vec(),
+        original[0..n / 2].to_vec(),
+        vec![],
+        original[0..1].to_vec(),
+        original[0..3].to_vec(),
+        original[0..n - 12].to_vec(),
+    ];
+
+    test_routine(original, modified);
+}
+
 #[test]
 fn test_vec_of_u64() {
     let original: Vec<u64> = vec![1, 2, 3, 4, 5];