Begin implementing cached hashing in types

eth2/utils/cached_tree_hash/src/impls.rs

@@ -1,12 +1,46 @@
 use super::*;
 use crate::merkleize::merkleize;
+use ethereum_types::H256;
 
 pub mod vec;
 
-impl CachedTreeHash<u64> for u64 {
+macro_rules! impl_for_single_leaf_int {
+    ($type: ident) => {
+        impl CachedTreeHash<$type> for $type {
+            fn new_tree_hash_cache(&self, _depth: usize) -> Result<TreeHashCache, Error> {
+                Ok(TreeHashCache::from_bytes(
+                    merkleize(self.to_le_bytes().to_vec()),
+                    false,
+                    None,
+                )?)
+            }
+
+            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.to_le_bytes().to_vec());
+                cache.maybe_update_chunk(cache.chunk_index, &leaf)?;
+
+                cache.chunk_index += 1;
+
+                Ok(())
+            }
+        }
+    };
+}
+
+impl_for_single_leaf_int!(u8);
+impl_for_single_leaf_int!(u16);
+impl_for_single_leaf_int!(u32);
+impl_for_single_leaf_int!(u64);
+impl_for_single_leaf_int!(usize);
+
+impl CachedTreeHash<bool> for bool {
     fn new_tree_hash_cache(&self, _depth: usize) -> Result<TreeHashCache, Error> {
         Ok(TreeHashCache::from_bytes(
-            merkleize(self.to_le_bytes().to_vec()),
+            merkleize((*self as u8).to_le_bytes().to_vec()),
             false,
             None,
         )?)
@@ -17,20 +51,19 @@ impl CachedTreeHash<u64> for u64 {
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {
-        let leaf = merkleize(self.to_le_bytes().to_vec());
+        let leaf = merkleize((*self as u8).to_le_bytes().to_vec());
         cache.maybe_update_chunk(cache.chunk_index, &leaf)?;
 
         cache.chunk_index += 1;
-        // cache.overlay_index += 1;
 
         Ok(())
     }
 }
 
-impl CachedTreeHash<usize> for usize {
+impl CachedTreeHash<[u8; 4]> for [u8; 4] {
     fn new_tree_hash_cache(&self, _depth: usize) -> Result<TreeHashCache, Error> {
         Ok(TreeHashCache::from_bytes(
-            merkleize(self.to_le_bytes().to_vec()),
+            merkleize(self.to_vec()),
             false,
             None,
         )?)
@@ -41,11 +74,33 @@ impl CachedTreeHash<usize> for usize {
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {
-        let leaf = merkleize(self.to_le_bytes().to_vec());
+        let leaf = merkleize(self.to_vec());
+        cache.maybe_update_chunk(cache.chunk_index, &leaf)?;
+
+        cache.chunk_index += 1;
+
+        Ok(())
+    }
+}
+
+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()),
+            false,
+            None,
+        )?)
+    }
+
+    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, &leaf)?;
 
         cache.chunk_index += 1;
-        // cache.overlay_index += 1;
 
         Ok(())
     }

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

@@ -95,6 +95,11 @@ pub fn update_tree_hash_cache<T: CachedTreeHash<T>>(
     let old_overlay = cache.get_overlay(cache.schema_index, cache.chunk_index)?;
     let new_overlay = BTreeOverlay::new(vec, cache.chunk_index, old_overlay.depth);
 
+    dbg!(cache.schema_index);
+    dbg!(cache.schemas.len());
+    dbg!(&old_overlay);
+    dbg!(&new_overlay);
+
     cache.replace_overlay(cache.schema_index, cache.chunk_index, new_overlay.clone())?;
 
     cache.schema_index += 1;

eth2/utils/cached_tree_hash/src/lib.rs

@@ -1,4 +1,5 @@
 use hashing::hash;
+use merkleize::num_unsanitized_leaves;
 use std::ops::Range;
 use tree_hash::{TreeHash, TreeHashType, BYTES_PER_CHUNK, HASHSIZE};
 
@@ -62,3 +63,92 @@ impl CachedTreeHasher {
         Ok(self.cache.root()?.to_vec())
     }
 }
+
+#[macro_export]
+macro_rules! cached_tree_hash_ssz_encoding_as_vector {
+    ($type: ident, $num_bytes: expr) => {
+        impl cached_tree_hash::CachedTreeHash<$type> for $type {
+            fn new_tree_hash_cache(
+                &self,
+                depth: usize,
+            ) -> Result<cached_tree_hash::TreeHashCache, cached_tree_hash::Error> {
+                let (cache, _schema) = cached_tree_hash::impls::vec::new_tree_hash_cache(
+                    &ssz::ssz_encode(self),
+                    depth,
+                )?;
+
+                Ok(cache)
+            }
+
+            fn tree_hash_cache_schema(&self, depth: usize) -> cached_tree_hash::BTreeSchema {
+                let lengths =
+                    vec![1; cached_tree_hash::merkleize::num_unsanitized_leaves($num_bytes)];
+                cached_tree_hash::BTreeSchema::from_lengths(depth, lengths)
+            }
+
+            fn update_tree_hash_cache(
+                &self,
+                cache: &mut cached_tree_hash::TreeHashCache,
+            ) -> Result<(), cached_tree_hash::Error> {
+                cached_tree_hash::impls::vec::update_tree_hash_cache(
+                    &ssz::ssz_encode(self),
+                    cache,
+                )?;
+
+                Ok(())
+            }
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! cached_tree_hash_bytes_as_list {
+    ($type: ident) => {
+        impl cached_tree_hash::CachedTreeHash<$type> for $type {
+            fn new_tree_hash_cache(
+                &self,
+                depth: usize,
+            ) -> Result<cached_tree_hash::TreeHashCache, cached_tree_hash::Error> {
+                let bytes = self.to_bytes();
+
+                let (mut cache, schema) =
+                    cached_tree_hash::impls::vec::new_tree_hash_cache(&bytes, depth)?;
+
+                cache.add_length_nodes(schema.into_overlay(0).chunk_range(), bytes.len())?;
+
+                Ok(cache)
+            }
+
+            fn num_tree_hash_cache_chunks(&self) -> usize {
+                // Add two extra nodes to cater for the node before and after to allow mixing-in length.
+                cached_tree_hash::BTreeOverlay::new(self, 0, 0).num_chunks() + 2
+            }
+
+            fn tree_hash_cache_schema(&self, depth: usize) -> cached_tree_hash::BTreeSchema {
+                cached_tree_hash::impls::vec::produce_schema(&ssz::ssz_encode(self), depth)
+            }
+
+            fn update_tree_hash_cache(
+                &self,
+                cache: &mut cached_tree_hash::TreeHashCache,
+            ) -> Result<(), cached_tree_hash::Error> {
+                let bytes = self.to_bytes();
+
+                // Skip the length-mixed-in root node.
+                cache.chunk_index += 1;
+
+                // Update the cache, returning the new overlay.
+                let new_overlay =
+                    cached_tree_hash::impls::vec::update_tree_hash_cache(&bytes, cache)?;
+
+                // Mix in length
+                cache.mix_in_length(new_overlay.chunk_range(), bytes.len())?;
+
+                // Skip an extra node to clear the length node.
+                cache.chunk_index = new_overlay.next_node() + 1;
+
+                Ok(())
+            }
+        }
+    };
+}

eth2/utils/cached_tree_hash/src/merkleize.rs

@@ -60,7 +60,7 @@ fn last_leaf_needs_padding(num_bytes: usize) -> bool {
 }
 
 /// Rounds up
-fn num_unsanitized_leaves(num_bytes: usize) -> usize {
+pub fn num_unsanitized_leaves(num_bytes: usize) -> usize {
     (num_bytes + HASHSIZE - 1) / HASHSIZE
 }