Refactor tree hashing (#861)

* Pre-allocated tree hash caches

* Add SmallVec to tree hash cache

* Avoid allocation for validator.pubkey

* Avoid iterator which seems to be doing heap alloc

* Add more smallvecs

* MOAR SMALLVEC

* Move non-test code to Hash256 tree hash

* Fix byte ordering error

* Add incomplete but working merkle stream impl

* Fix zero hash error

* Add zero hash fn

* Add MerkleStream comments

* Add smallvec, tidy

* Integrate into tree hash derive

* Update ssz_types tree hash

* Don't heap alloc for mix in length

* Add byte-level streaming to MerkleStream

* Avoid recursion in write method

* Update BLS to MerkleStream

* Fix some not-compiling tests

* Remove debug profiling

* Remove code duplication

* Move beacon state tree hash to new hasher

* Fix failing tests

* Update comments

* Add some fast-paths to tree_hash::merkle_root

* Remove unncessary test

* Rename MerkleStream -> MerkleHasher

* Rename new_with_leaf_count -> with_leaves

* Tidy

* Remove NonZeroUsize

* Remove todo

* Update smallvec

eth2/utils/cached_tree_hash/src/cache.rs

@@ -1,6 +1,8 @@
 use crate::cache_arena;
+use crate::SmallVec8;
 use crate::{Error, Hash256};
 use eth2_hashing::{hash32_concat, ZERO_HASHES};
+use smallvec::smallvec;
 use ssz_derive::{Decode, Encode};
 use tree_hash::BYTES_PER_CHUNK;
 
@@ -17,28 +19,25 @@ pub struct TreeHashCache {
     ///
     /// The leaves are contained in `self.layers[self.depth]`, and each other layer `i`
     /// contains the parents of the nodes in layer `i + 1`.
-    layers: Vec<CacheArenaAllocation>,
+    layers: SmallVec8<CacheArenaAllocation>,
 }
 
 impl TreeHashCache {
     /// Create a new cache with the given `depth` with enough nodes allocated to suit `leaves`. All
     /// leaves are set to `Hash256::zero()`.
     pub fn new(arena: &mut CacheArena, depth: usize, leaves: usize) -> Self {
-        // TODO: what about when leaves is zero?
-        let layers = (0..=depth)
-            .map(|i| {
-                let vec = arena.alloc();
-                vec.extend_with_vec(
-                    arena,
-                    vec![Hash256::zero(); nodes_per_layer(i, depth, leaves)],
-                )
-                .expect(
-                    "A newly allocated sub-arena cannot fail unless it has reached max capacity",
-                );
+        let mut layers = SmallVec8::with_capacity(depth + 1);
 
-                vec
-            })
-            .collect();
+        for i in 0..=depth {
+            let vec = arena.alloc();
+            vec.extend_with_vec(
+                arena,
+                smallvec![Hash256::zero(); nodes_per_layer(i, depth, leaves)],
+            )
+            .expect("A newly allocated sub-arena cannot fail unless it has reached max capacity");
+
+            layers.push(vec)
+        }
 
         TreeHashCache {
             initialized: false,
@@ -62,7 +61,7 @@ impl TreeHashCache {
         &mut self,
         arena: &mut CacheArena,
         mut leaves: impl Iterator<Item = [u8; BYTES_PER_CHUNK]> + ExactSizeIterator,
-    ) -> Result<Vec<usize>, Error> {
+    ) -> Result<SmallVec8<usize>, Error> {
         let new_leaf_count = leaves.len();
 
         if new_leaf_count < self.leaves().len(arena)? {
@@ -71,21 +70,19 @@ impl TreeHashCache {
             return Err(Error::TooManyLeaves);
         }
 
+        let mut dirty = SmallVec8::new();
+
         // Update the existing leaves
-        let mut dirty = self
-            .leaves()
+        self.leaves()
             .iter_mut(arena)?
             .enumerate()
             .zip(&mut leaves)
-            .flat_map(|((i, leaf), new_leaf)| {
+            .for_each(|((i, leaf), new_leaf)| {
                 if !self.initialized || leaf.as_bytes() != new_leaf {
                     leaf.assign_from_slice(&new_leaf);
-                    Some(i)
-                } else {
-                    None
+                    dirty.push(i);
                 }
-            })
-            .collect::<Vec<_>>();
+            });
 
         // Push the rest of the new leaves (if any)
         dirty.extend(self.leaves().len(arena)?..new_leaf_count);
@@ -101,7 +98,7 @@ impl TreeHashCache {
     pub fn update_merkle_root(
         &mut self,
         arena: &mut CacheArena,
-        mut dirty_indices: Vec<usize>,
+        mut dirty_indices: SmallVec8<usize>,
     ) -> Result<Hash256, Error> {
         if dirty_indices.is_empty() {
             return Ok(self.root(arena));
@@ -164,8 +161,13 @@ impl TreeHashCache {
 }
 
 /// Compute the dirty indices for one layer up.
-fn lift_dirty(dirty_indices: &[usize]) -> Vec<usize> {
-    let mut new_dirty = dirty_indices.iter().map(|i| *i / 2).collect::<Vec<_>>();
+fn lift_dirty(dirty_indices: &[usize]) -> SmallVec8<usize> {
+    let mut new_dirty = SmallVec8::with_capacity(dirty_indices.len());
+
+    for i in 0..dirty_indices.len() {
+        new_dirty.push(dirty_indices[i] / 2)
+    }
+
     new_dirty.dedup();
     new_dirty
 }
@@ -202,6 +204,21 @@ fn nodes_per_layer(layer: usize, depth: usize, leaves: usize) -> usize {
 mod test {
     use super::*;
 
+    #[test]
+    fn zero_leaves() {
+        let arena = &mut CacheArena::default();
+
+        let depth = 3;
+        let num_leaves = 0;
+
+        let mut cache = TreeHashCache::new(arena, depth, num_leaves);
+        let leaves: Vec<[u8; BYTES_PER_CHUNK]> = vec![];
+
+        cache
+            .recalculate_merkle_root(arena, leaves.into_iter())
+            .expect("should calculate root");
+    }
+
     #[test]
     fn test_node_per_layer_unbalanced_tree() {
         assert_eq!(nodes_per_layer(0, 3, 5), 1);

eth2/utils/cached_tree_hash/src/cache_arena.rs

@@ -1,3 +1,4 @@
+use crate::SmallVec8;
 use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use std::marker::PhantomData;
@@ -27,6 +28,14 @@ pub struct CacheArena<T: Encode + Decode> {
 }
 
 impl<T: Encode + Decode> CacheArena<T> {
+    /// Instantiate self with a backing array of the given `capacity`.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            backing: Vec::with_capacity(capacity),
+            offsets: vec![],
+        }
+    }
+
     /// Produce an allocation of zero length at the end of the backing array.
     pub fn alloc(&mut self) -> CacheArenaAllocation<T> {
         let alloc_id = self.offsets.len();
@@ -204,7 +213,11 @@ pub struct CacheArenaAllocation<T> {
 
 impl<T: Encode + Decode> CacheArenaAllocation<T> {
     /// Grow the allocation in `arena`, appending `vec` to the current values.
-    pub fn extend_with_vec(&self, arena: &mut CacheArena<T>, vec: Vec<T>) -> Result<(), Error> {
+    pub fn extend_with_vec(
+        &self,
+        arena: &mut CacheArena<T>,
+        vec: SmallVec8<T>,
+    ) -> Result<(), Error> {
         let len = arena.len(self.alloc_id)?;
         arena.splice_forgetful(self.alloc_id, len..len, vec)?;
         Ok(())
@@ -264,6 +277,7 @@ impl<T: Encode + Decode> CacheArenaAllocation<T> {
 #[cfg(test)]
 mod tests {
     use crate::Hash256;
+    use smallvec::smallvec;
 
     type CacheArena = super::CacheArena<Hash256>;
     type CacheArenaAllocation = super::CacheArenaAllocation<Hash256>;
@@ -300,7 +314,7 @@ mod tests {
             len
         );
 
-        sub.extend_with_vec(arena, vec![hash(len), hash(len + 1)])
+        sub.extend_with_vec(arena, smallvec![hash(len), hash(len + 1)])
             .expect("should extend with vec");
         len += 2;
 

eth2/utils/cached_tree_hash/src/impls.rs

@@ -97,12 +97,10 @@ impl<N: Unsigned> CachedTreeHash<TreeHashCache> for VariableList<Hash256, N> {
         arena: &mut CacheArena,
         cache: &mut TreeHashCache,
     ) -> Result<Hash256, Error> {
-        Ok(Hash256::from_slice(&mix_in_length(
-            cache
-                .recalculate_merkle_root(arena, hash256_iter(&self))?
-                .as_bytes(),
+        Ok(mix_in_length(
+            &cache.recalculate_merkle_root(arena, hash256_iter(&self))?,
             self.len(),
-        )))
+        ))
     }
 }
 
@@ -121,12 +119,10 @@ impl<N: Unsigned> CachedTreeHash<TreeHashCache> for VariableList<u64, N> {
         arena: &mut CacheArena,
         cache: &mut TreeHashCache,
     ) -> Result<Hash256, Error> {
-        Ok(Hash256::from_slice(&mix_in_length(
-            cache
-                .recalculate_merkle_root(arena, u64_iter(&self))?
-                .as_bytes(),
+        Ok(mix_in_length(
+            &cache.recalculate_merkle_root(arena, u64_iter(&self))?,
             self.len(),
-        )))
+        ))
     }
 }
 

eth2/utils/cached_tree_hash/src/lib.rs

@@ -3,7 +3,9 @@ mod cache_arena;
 mod impls;
 #[cfg(test)]
 mod test;
+use smallvec::SmallVec;
 
+type SmallVec8<T> = SmallVec<[T; 8]>;
 pub type CacheArena = cache_arena::CacheArena<Hash256>;
 
 pub use crate::cache::TreeHashCache;

eth2/utils/cached_tree_hash/src/test.rs

@@ -76,7 +76,7 @@ fn fixed_vector_hash256() {
     let mut cache = vec.new_tree_hash_cache(arena);
 
     assert_eq!(
-        Hash256::from_slice(&vec.tree_hash_root()),
+        vec.tree_hash_root(),
         vec.recalculate_tree_hash_root(arena, &mut cache).unwrap()
     );
 }
@@ -90,7 +90,7 @@ fn fixed_vector_u64() {
     let mut cache = vec.new_tree_hash_cache(arena);
 
     assert_eq!(
-        Hash256::from_slice(&vec.tree_hash_root()),
+        vec.tree_hash_root(),
         vec.recalculate_tree_hash_root(arena, &mut cache).unwrap()
     );
 }
@@ -104,7 +104,7 @@ fn variable_list_hash256() {
     let mut cache = list.new_tree_hash_cache(arena);
 
     assert_eq!(
-        Hash256::from_slice(&list.tree_hash_root()),
+        list.tree_hash_root(),
         list.recalculate_tree_hash_root(arena, &mut cache).unwrap()
     );
 }