Allow for building cached vec

eth2/utils/ssz/src/cached_tree_hash.rs

@@ -22,9 +22,18 @@ pub enum Error {
     NoChildrenForHashing((usize, usize)),
 }
 
+#[derive(Debug, PartialEq, Clone)]
+pub enum ItemType {
+    Basic,
+    List,
+    Composite,
+}
+
 pub trait CachedTreeHash {
     type Item: CachedTreeHash;
 
+    fn item_type() -> ItemType;
+
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error>;
 
     /// Return the number of bytes when this element is encoded as raw SSZ _without_ length
@@ -35,6 +44,8 @@ pub trait CachedTreeHash {
 
     fn num_child_nodes(&self) -> usize;
 
+    fn packed_encoding(&self) -> Vec<u8>;
+
     fn cached_hash_tree_root(
         &self,
         other: &Self::Item,
@@ -101,6 +112,7 @@ impl TreeHashCache {
         // Merkleize the leaves, then split the leaf nodes off them. Then, replace all-zeros
         // internal nodes created earlier with the internal nodes generated by `merkleize`.
         let mut merkleized = merkleize(leaves);
+        dbg!(&merkleized);
         merkleized.split_off(internal_node_bytes);
         cache.splice(0..internal_node_bytes, merkleized);
 

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

@@ -4,6 +4,10 @@ use crate::{ssz_encode, Encodable};
 impl CachedTreeHash for u64 {
     type Item = Self;
 
+    fn item_type() -> ItemType {
+        ItemType::Basic
+    }
+
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
         Ok(TreeHashCache::from_bytes(merkleize(ssz_encode(self)))?)
     }
@@ -20,6 +24,10 @@ impl CachedTreeHash for u64 {
         0
     }
 
+    fn packed_encoding(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+
     fn cached_hash_tree_root(
         &self,
         other: &Self,
@@ -35,38 +43,73 @@ impl CachedTreeHash for u64 {
     }
 }
 
-/*
 impl<T> CachedTreeHash for Vec<T>
 where
-    T: CachedTreeHash + Encodable,
+    T: CachedTreeHash,
 {
     type Item = Self;
 
-    fn build_cache_bytes(&self) -> Vec<u8> {
-        let num_packed_bytes = self.num_bytes();
-        let num_leaves = num_sanitized_leaves(num_packed_bytes);
+    fn item_type() -> ItemType {
+        ItemType::List
+    }
 
-        let mut packed = Vec::with_capacity(num_leaves * HASHSIZE);
+    fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
+        match T::item_type() {
+            ItemType::Basic => {
+                let num_packed_bytes = self.num_bytes();
+                let num_leaves = num_sanitized_leaves(num_packed_bytes);
+
+                let mut packed = Vec::with_capacity(num_leaves * HASHSIZE);
+
+                for item in self {
+                    packed.append(&mut item.packed_encoding());
+                }
+
+                let packed = sanitise_bytes(packed);
+
+                TreeHashCache::from_bytes(merkleize(packed))
+            }
+            ItemType::Composite | ItemType::List => {
+                let subtrees = self
+                    .iter()
+                    .map(|item| TreeHashCache::new(item))
+                    .collect::<Result<Vec<TreeHashCache>, _>>()?;
+
+                TreeHashCache::from_leaves_and_subtrees(self, subtrees)
+            }
+        }
+    }
+
+    fn offsets(&self) -> Result<Vec<usize>, Error> {
+        let mut offsets = vec![];
 
         for item in self {
-            packed.append(&mut ssz_encode(item));
+            offsets.push(item.offsets()?.iter().sum())
         }
 
-        let packed = sanitise_bytes(packed);
+        Ok(offsets)
+    }
 
-        merkleize(packed)
+    fn num_child_nodes(&self) -> usize {
+        // TODO
+        42
     }
 
     fn num_bytes(&self) -> usize {
         self.iter().fold(0, |acc, item| acc + item.num_bytes())
     }
 
+    fn packed_encoding(&self) -> Vec<u8> {
+        panic!("List should never be packed")
+    }
+
     fn cached_hash_tree_root(
         &self,
         other: &Self::Item,
         cache: &mut TreeHashCache,
         chunk: usize,
-    ) -> Option<usize> {
+    ) -> Result<usize, Error> {
+        /*
         let num_packed_bytes = self.num_bytes();
         let num_leaves = num_sanitized_leaves(num_packed_bytes);
 
@@ -103,6 +146,17 @@ where
         }
 
         Some(chunk + num_nodes)
+        */
+        // TODO
+        Ok(42)
     }
 }
+
+/*
+fn get_packed_leaves<T>(vec: Vec<T>) -> Vec<u8>
+where
+    T: Encodable,
+{
+    //
+}
 */

eth2/utils/ssz/src/cached_tree_hash/tests.rs

@@ -1,5 +1,6 @@
 use super::*;
-use int_to_bytes::int_to_bytes32;
+use crate::Encodable;
+use int_to_bytes::{int_to_bytes32, int_to_bytes8};
 
 #[derive(Clone)]
 pub struct Inner {
@@ -12,6 +13,10 @@ pub struct Inner {
 impl CachedTreeHash for Inner {
     type Item = Self;
 
+    fn item_type() -> ItemType {
+        ItemType::Composite
+    }
+
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
         let tree = TreeHashCache::from_leaves_and_subtrees(
             self,
@@ -60,6 +65,10 @@ impl CachedTreeHash for Inner {
         num_nodes(leaves) + children - 1
     }
 
+    fn packed_encoding(&self) -> Vec<u8> {
+        panic!("Struct should never be packed")
+    }
+
     fn cached_hash_tree_root(
         &self,
         other: &Self,
@@ -97,6 +106,10 @@ pub struct Outer {
 impl CachedTreeHash for Outer {
     type Item = Self;
 
+    fn item_type() -> ItemType {
+        ItemType::Composite
+    }
+
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
         let tree = TreeHashCache::from_leaves_and_subtrees(
             self,
@@ -139,6 +152,10 @@ impl CachedTreeHash for Outer {
         Ok(offsets)
     }
 
+    fn packed_encoding(&self) -> Vec<u8> {
+        panic!("Struct should never be packed")
+    }
+
     fn cached_hash_tree_root(
         &self,
         other: &Self,
@@ -371,6 +388,64 @@ fn large_vec_of_u64_builds() {
 
     assert_eq!(expected, cache);
 }
+*/
+
+#[test]
+fn vec_of_inner_builds() {
+    let numbers: Vec<u64> = (0..12).collect();
+
+    let mut leaves = vec![];
+    let mut full_bytes = vec![];
+
+    for n in numbers.chunks(4) {
+        let mut merkle = merkleize(join(vec![
+            int_to_bytes32(n[0]),
+            int_to_bytes32(n[1]),
+            int_to_bytes32(n[2]),
+            int_to_bytes32(n[3]),
+        ]));
+        leaves.append(&mut merkle[0..HASHSIZE].to_vec());
+        full_bytes.append(&mut merkle);
+    }
+
+    let mut expected = merkleize(leaves);
+    expected.splice(3 * HASHSIZE.., full_bytes);
+    expected.append(&mut vec![0; HASHSIZE]);
+
+    let my_vec = vec![
+        Inner {
+            a: 0,
+            b: 1,
+            c: 2,
+            d: 3,
+        },
+        Inner {
+            a: 4,
+            b: 5,
+            c: 6,
+            d: 7,
+        },
+        Inner {
+            a: 8,
+            b: 9,
+            c: 10,
+            d: 11,
+        },
+    ];
+
+    let cache: Vec<u8> = TreeHashCache::new(&my_vec).unwrap().into();
+
+    assert_trees_eq(&expected, &cache);
+}
+
+/// Provides detailed assertions when comparing merkle trees.
+fn assert_trees_eq(a: &[u8], b: &[u8]) {
+    assert_eq!(a.len(), b.len(), "Byte lens different");
+    for i in 0..a.len() / HASHSIZE {
+        let range = i * HASHSIZE..(i + 1) * HASHSIZE;
+        assert_eq!(a[range.clone()], b[range], "Chunk {} different", i);
+    }
+}
 
 #[test]
 fn vec_of_u64_builds() {
@@ -387,11 +462,10 @@ fn vec_of_u64_builds() {
 
     let my_vec = vec![1, 2, 3, 4, 5];
 
-    let cache = my_vec.build_cache_bytes();
+    let cache: Vec<u8> = TreeHashCache::new(&my_vec).unwrap().into();
 
     assert_eq!(expected, cache);
 }
-*/
 
 #[test]
 fn merkleize_odd() {