Update CachedTreeHash trait to return schema

eth2/types/src/tree_hash_vector.rs

@@ -68,18 +68,8 @@ where
         Ok(cache)
     }
 
-    fn num_tree_hash_cache_chunks(&self) -> usize {
+    fn tree_hash_cache_schema(&self, depth: usize) -> cached_tree_hash::BTreeSchema {
-        cached_tree_hash::BTreeOverlay::new(self, 0, 0)
+        cached_tree_hash::impls::vec::produce_schema(self, depth)
-            .and_then(|o| Ok(o.num_chunks()))
-            .unwrap_or_else(|_| 1)
-    }
-
-    fn tree_hash_cache_overlay(
-        &self,
-        chunk_offset: usize,
-        depth: usize,
-    ) -> Result<cached_tree_hash::BTreeOverlay, cached_tree_hash::Error> {
-        cached_tree_hash::impls::vec::produce_overlay(self, chunk_offset, depth)
     }
 
     fn update_tree_hash_cache(

eth2/utils/cached_tree_hash/src/btree_overlay.rs

@@ -1,10 +1,39 @@
 use super::*;
 
 #[derive(Debug, PartialEq, Clone)]
-pub struct BTreeOverlay {
+pub struct BTreeSchema {
-    pub offset: usize,
     pub depth: usize,
-    pub lengths: Vec<usize>,
+    lengths: Vec<usize>,
+}
+
+impl BTreeSchema {
+    pub fn from_lengths(depth: usize, lengths: Vec<usize>) -> Self {
+        Self { depth, lengths }
+    }
+
+    pub fn into_overlay(self, offset: usize) -> BTreeOverlay {
+        BTreeOverlay {
+            offset,
+            depth: self.depth,
+            lengths: self.lengths,
+        }
+    }
+}
+
+impl Into<BTreeSchema> for BTreeOverlay {
+    fn into(self) -> BTreeSchema {
+        BTreeSchema {
+            depth: self.depth,
+            lengths: self.lengths,
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Clone)]
+pub struct BTreeOverlay {
+    offset: usize,
+    pub depth: usize,
+    lengths: Vec<usize>,
 }
 
 impl BTreeOverlay {
@@ -12,15 +41,8 @@ impl BTreeOverlay {
     where
         T: CachedTreeHash<T>,
     {
-        item.tree_hash_cache_overlay(initial_offset, depth)
+        item.tree_hash_cache_schema(depth)
-    }
+            .into_overlay(initial_offset)
-
-    pub fn from_lengths(offset: usize, depth: usize, lengths: Vec<usize>) -> Self {
-        Self {
-            offset,
-            depth,
-            lengths,
-        }
     }
 
     pub fn num_leaf_nodes(&self) -> usize {
@@ -166,7 +188,7 @@ mod test {
     use super::*;
 
     fn get_tree_a(n: usize) -> BTreeOverlay {
-        BTreeOverlay::from_lengths(0, 0, vec![1; n])
+        BTreeSchema::from_lengths(0, vec![1; n]).into_overlay(0)
     }
 
     #[test]
@@ -204,7 +226,7 @@ mod test {
         let tree = get_tree_a(2);
         assert_eq!(tree.chunk_range(), 0..3);
 
-        let tree = BTreeOverlay::from_lengths(11, 0, vec![1, 1]);
+        let tree = BTreeSchema::from_lengths(0, vec![1, 1]).into_overlay(11);
         assert_eq!(tree.chunk_range(), 11..14);
     }
 

eth2/utils/cached_tree_hash/src/impls.rs

@@ -12,12 +12,8 @@ impl CachedTreeHash<u64> for u64 {
         )?)
     }
 
-    fn num_tree_hash_cache_chunks(&self) -> usize {
+    fn tree_hash_cache_schema(&self, depth: usize) -> BTreeSchema {
-        1
+        BTreeSchema::from_lengths(depth, vec![1])
-    }
-
-    fn tree_hash_cache_overlay(&self, _chunk_offset: usize, _depth: usize) -> BTreeOverlay {
-        unreachable!("Basic should not produce overlay");
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {
@@ -40,12 +36,8 @@ impl CachedTreeHash<usize> for usize {
         )?)
     }
 
-    fn num_tree_hash_cache_chunks(&self) -> usize {
+    fn tree_hash_cache_schema(&self, depth: usize) -> BTreeSchema {
-        1
+        BTreeSchema::from_lengths(depth, vec![1])
-    }
-
-    fn tree_hash_cache_overlay(&self, _chunk_offset: usize, _depth: usize) -> BTreeOverlay {
-        unreachable!("Basic should not produce overlay");
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {

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

@@ -6,9 +6,9 @@ where
     T: CachedTreeHash<T> + TreeHash,
 {
     fn new_tree_hash_cache(&self, depth: usize) -> Result<TreeHashCache, Error> {
-        let (mut cache, overlay) = new_tree_hash_cache(self, depth)?;
+        let (mut cache, schema) = new_tree_hash_cache(self, depth)?;
 
-        cache.add_length_nodes(overlay.chunk_range(), self.len())?;
+        cache.add_length_nodes(schema.into_overlay(0).chunk_range(), self.len())?;
 
         Ok(cache)
     }
@@ -18,8 +18,8 @@ where
         BTreeOverlay::new(self, 0, 0).num_chunks() + 2
     }
 
-    fn tree_hash_cache_overlay(&self, chunk_offset: usize, depth: usize) -> BTreeOverlay {
+    fn tree_hash_cache_schema(&self, depth: usize) -> BTreeSchema {
-        produce_overlay(self, chunk_offset, depth)
+        produce_schema(self, depth)
     }
 
     fn update_tree_hash_cache(&self, cache: &mut TreeHashCache) -> Result<(), Error> {
@@ -42,14 +42,14 @@ where
 pub fn new_tree_hash_cache<T: CachedTreeHash<T>>(
     vec: &Vec<T>,
     depth: usize,
-) -> Result<(TreeHashCache, BTreeOverlay), Error> {
+) -> Result<(TreeHashCache, BTreeSchema), Error> {
-    let overlay = vec.tree_hash_cache_overlay(0, depth);
+    let schema = vec.tree_hash_cache_schema(depth);
 
     let cache = match T::tree_hash_type() {
         TreeHashType::Basic => TreeHashCache::from_bytes(
             merkleize(get_packed_leaves(vec)?),
             false,
-            Some(overlay.clone()),
+            Some(schema.clone()),
         ),
         TreeHashType::Container | TreeHashType::List | TreeHashType::Vector => {
             let subtrees = vec
@@ -61,14 +61,10 @@ pub fn new_tree_hash_cache<T: CachedTreeHash<T>>(
         }
     }?;
 
-    Ok((cache, overlay))
+    Ok((cache, schema))
 }
 
-pub fn produce_overlay<T: CachedTreeHash<T>>(
+pub fn produce_schema<T: CachedTreeHash<T>>(vec: &Vec<T>, depth: usize) -> BTreeSchema {
-    vec: &Vec<T>,
-    chunk_offset: usize,
-    depth: usize,
-) -> BTreeOverlay {
     let lengths = match T::tree_hash_type() {
         TreeHashType::Basic => {
             // Ceil division.
@@ -89,7 +85,7 @@ pub fn produce_overlay<T: CachedTreeHash<T>>(
         }
     };
 
-    BTreeOverlay::from_lengths(chunk_offset, depth, lengths)
+    BTreeSchema::from_lengths(depth, lengths)
 }
 
 pub fn update_tree_hash_cache<T: CachedTreeHash<T>>(

eth2/utils/cached_tree_hash/src/lib.rs

@@ -9,14 +9,16 @@ pub mod merkleize;
 mod resize;
 mod tree_hash_cache;
 
-pub use btree_overlay::BTreeOverlay;
+pub use btree_overlay::{BTreeOverlay, BTreeSchema};
 pub use errors::Error;
 pub use tree_hash_cache::TreeHashCache;
 
 pub trait CachedTreeHash<Item>: TreeHash {
-    fn tree_hash_cache_overlay(&self, chunk_offset: usize, depth: usize) -> BTreeOverlay;
+    fn tree_hash_cache_schema(&self, depth: usize) -> BTreeSchema;
 
-    fn num_tree_hash_cache_chunks(&self) -> usize;
+    fn num_tree_hash_cache_chunks(&self) -> usize {
+        self.tree_hash_cache_schema(0).into_overlay(0).num_chunks()
+    }
 
     fn new_tree_hash_cache(&self, depth: usize) -> Result<TreeHashCache, Error>;
 

eth2/utils/cached_tree_hash/src/tree_hash_cache.rs

@@ -2,31 +2,6 @@ use super::*;
 use crate::merkleize::{merkleize, pad_for_leaf_count};
 use int_to_bytes::int_to_bytes32;
 
-#[derive(Debug, PartialEq, Clone)]
-pub struct BTreeSchema {
-    pub depth: usize,
-    pub lengths: Vec<usize>,
-}
-
-impl BTreeSchema {
-    pub fn into_overlay(self, offset: usize) -> BTreeOverlay {
-        BTreeOverlay {
-            offset,
-            depth: self.depth,
-            lengths: self.lengths,
-        }
-    }
-}
-
-impl Into<BTreeSchema> for BTreeOverlay {
-    fn into(self) -> BTreeSchema {
-        BTreeSchema {
-            depth: self.depth,
-            lengths: self.lengths,
-        }
-    }
-}
-
 #[derive(Debug, PartialEq, Clone)]
 pub struct TreeHashCache {
     pub cache: Vec<u8>,
@@ -113,14 +88,14 @@ impl TreeHashCache {
     pub fn from_bytes(
         bytes: Vec<u8>,
         initial_modified_state: bool,
-        overlay: Option<BTreeOverlay>,
+        schema: Option<BTreeSchema>,
     ) -> Result<Self, Error> {
         if bytes.len() % BYTES_PER_CHUNK > 0 {
             return Err(Error::BytesAreNotEvenChunks(bytes.len()));
         }
 
-        let schemas = match overlay {
+        let schemas = match schema {
-            Some(overlay) => vec![overlay.into()],
+            Some(schema) => vec![schema],
             None => vec![],
         };
 

eth2/utils/tree_hash_derive/src/lib.rs

@@ -75,14 +75,14 @@ pub fn subtree_derive(input: TokenStream) -> TokenStream {
                 cached_tree_hash::BTreeOverlay::new(self, 0, 0).num_chunks()
             }
 
-            fn tree_hash_cache_overlay(&self, chunk_offset: usize, depth: usize) ->cached_tree_hash::BTreeOverlay {
+            fn tree_hash_cache_schema(&self, depth: usize) -> cached_tree_hash::BTreeSchema {
                 let mut lengths = vec![];
 
                 #(
                     lengths.push(self.#idents_b.num_tree_hash_cache_chunks());
                 )*
 
-                cached_tree_hash::BTreeOverlay::from_lengths(chunk_offset, depth, lengths)
+                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> {