Refactor with TreeHashCache struct

eth2/utils/ssz/Cargo.toml

@@ -8,3 +8,4 @@ edition = "2018"
 bytes = "0.4.9"
 ethereum-types = "0.5"
 hashing = { path = "../hashing" }
+int_to_bytes = { path = "../int_to_bytes" }

eth2/utils/ssz/src/cached_tree_hash.rs

@@ -1,38 +1,97 @@
 use crate::ssz_encode;
 use hashing::hash;
+use int_to_bytes::int_to_bytes32;
 
 const BYTES_PER_CHUNK: usize = 32;
 const HASHSIZE: usize = 32;
 const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;
 
+pub struct TreeHashCache<'a> {
+    chunk_offset: usize,
+    cache: &'a mut [u8],
+    chunk_modified: &'a mut [bool],
+}
+
+impl<'a> TreeHashCache<'a> {
+    pub fn increment(&mut self) {
+        self.chunk_offset += 1
+    }
+
+    pub fn modify_current_chunk(&mut self, to: &[u8]) -> Option<()> {
+        self.modify_chunk(0, to)
+    }
+
+    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> {
+        let start = chunk * BYTES_PER_CHUNK;
+        let end = start + BYTES_PER_CHUNK;
+        self.cache.get_mut(start..end)?.copy_from_slice(to);
+
+        self.chunk_modified[chunk] = true;
+
+        Some(())
+    }
+
+    pub fn changed(&self, chunk: usize) -> Option<bool> {
+        self.chunk_modified.get(chunk).cloned()
+    }
+
+    pub fn children_modified(&self, parent_chunk: usize) -> Option<bool> {
+        let children = children(parent_chunk);
+
+        Some(self.changed(children.0)? | self.changed(children.1)?)
+    }
+
+    pub fn hash_children(&self, parent_chunk: usize) -> Option<Vec<u8>> {
+        let children = children(parent_chunk);
+
+        let start = children.0 * BYTES_PER_CHUNK;
+        let end = start + BYTES_PER_CHUNK * 2;
+
+        Some(hash(&self.cache.get(start..end)?))
+    }
+
+    pub fn just_the_leaves(&mut self, leaves: usize) -> Option<TreeHashCache> {
+        let nodes = num_nodes(leaves);
+        let internal = nodes - leaves;
+        let leaves_start = (self.chunk_offset + internal) * HASHSIZE;
+
+        Some(TreeHashCache {
+            chunk_offset: self.chunk_offset + internal,
+            cache: self.cache.get_mut(leaves_start..leaves * HASHSIZE)?,
+            chunk_modified: self
+                .chunk_modified
+                .get_mut(self.chunk_offset..self.chunk_offset + leaves)?,
+        })
+    }
+}
+
+fn children(parent: usize) -> (usize, usize) {
+    ((2 * parent + 1), (2 * parent + 2))
+}
+
+fn num_nodes(num_leaves: usize) -> usize {
+    2 * num_leaves - 1
+}
+
 pub trait CachedTreeHash {
-    fn cached_hash_tree_root(
+    fn build_cache_bytes(&self) -> Vec<u8>;
-        &self,
+
-        other: &Self,
+    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>;
-        cache: &mut [u8],
-        i: usize,
-        changes: Vec<bool>,
-    ) -> Option<(usize, Vec<bool>)>;
 }
 
 impl CachedTreeHash for u64 {
-    fn cached_hash_tree_root(
+    fn build_cache_bytes(&self) -> Vec<u8> {
-        &self,
+        merkleize(&ssz_encode(self))
-        other: &Self,
+    }
-        cache: &mut [u8],
-        i: usize,
-        mut changes: Vec<bool>,
-    ) -> Option<(usize, Vec<bool>)> {
-        if self != other {
-            cache
-                .get_mut(i..i + HASHSIZE)?
-                .copy_from_slice(&mut hash(&ssz_encode(self)));
-            changes.push(true);
-        } else {
-            changes.push(false);
-        };
 
-        Some((i + HASHSIZE, changes))
+    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
+        if self != other {
+            cache.modify_current_chunk(&hash(&ssz_encode(self)));
+        }
+
+        cache.increment();
+
+        Some(())
     }
 }
 
@@ -44,70 +103,137 @@ pub struct Inner {
 }
 
 impl CachedTreeHash for Inner {
-    fn cached_hash_tree_root(
+    fn build_cache_bytes(&self) -> Vec<u8> {
-        &self,
+        let mut leaves = vec![];
-        other: &Self,
-        cache: &mut [u8],
-        i: usize,
-        mut changes: Vec<bool>,
-    ) -> Option<(usize, Vec<bool>)> {
-        let original_start = i;
 
-        let leaves = 4;
+        leaves.append(&mut self.a.build_cache_bytes());
-        let nodes = num_nodes(leaves);
+        leaves.append(&mut self.b.build_cache_bytes());
-        let internal = nodes - leaves;
+        leaves.append(&mut self.c.build_cache_bytes());
-        let leaves_start = i + internal * HASHSIZE;
+        leaves.append(&mut self.d.build_cache_bytes());
 
-        let mut leaf_changes = {
+        merkleize(&leaves)
-            let leaf_changes = Vec::with_capacity(leaves);
+    }
-            let leaf_start = leaves_start;
 
-            let (leaf_start, leaf_changes) =
+    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
-                self.a
+        let num_leaves = 4;
-                    .cached_hash_tree_root(&other.a, cache, leaf_start, leaf_changes)?;
-            let (leaf_start, leaf_changes) =
-                self.b
-                    .cached_hash_tree_root(&other.b, cache, leaf_start, leaf_changes)?;
-            let (leaf_start, leaf_changes) =
-                self.c
-                    .cached_hash_tree_root(&other.c, cache, leaf_start, leaf_changes)?;
-            let (_leaf_start, leaf_changes) =
-                self.d
-                    .cached_hash_tree_root(&other.d, cache, leaf_start, leaf_changes)?;
 
-            leaf_changes
+        let mut leaf_cache = cache.just_the_leaves(num_leaves)?;
-        };
+        self.a.cached_hash_tree_root(&other.a, &mut leaf_cache)?;
+        self.b.cached_hash_tree_root(&other.b, &mut leaf_cache)?;
+        self.c.cached_hash_tree_root(&other.c, &mut leaf_cache)?;
+        self.d.cached_hash_tree_root(&other.d, &mut leaf_cache)?;
 
-        let any_changes = leaf_changes.iter().any(|&c| c);
+        let nodes = num_nodes(num_leaves);
+        let internal_chunks = nodes - num_leaves;
 
-        changes.resize(changes.len() + internal, false);
+        for chunk in 0..internal_chunks {
-        changes.append(&mut leaf_changes);
+            if cache.children_modified(chunk)? {
-
+                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
-        if any_changes {
-            let mut i = internal;
-
-            while i > 0 {
-                let children = children(i);
-
-                if changes[children.0] | changes[children.1] {
-                    changes[parent(i)] = true;
-
-                    let children_start = children.0 * HASHSIZE;
-                    let children_end = children_start + 2 * HASHSIZE;
-                    let hash = hash(&cache.get(children_start..children_end)?);
-
-                    cache
-                        .get_mut(i * HASHSIZE..(i + 1) * HASHSIZE)?
-                        .copy_from_slice(&hash);
-                }
-                i += 1
             }
         }
 
-        Some((42, vec![any_changes]))
+        Some(())
     }
 }
 
+/// A reference function to test against.
+pub fn merkleize(values: &[u8]) -> Vec<u8> {
+    let leaves = values.len() / HASHSIZE;
+
+    if leaves == 0 || !leaves.is_power_of_two() {
+        panic!("Handle bad leaf count");
+    }
+
+    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
+    o.append(&mut values.to_vec());
+
+    let mut i = o.len();
+    let mut j = o.len() - values.len();
+
+    while i >= MERKLE_HASH_CHUNCK {
+        i -= MERKLE_HASH_CHUNCK;
+        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);
+
+        j -= HASHSIZE;
+        o[j..j + HASHSIZE].copy_from_slice(&hash);
+    }
+
+    o
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn join(many: Vec<&[u8]>) -> Vec<u8> {
+        let mut all = vec![];
+        for one in many {
+            all.extend_from_slice(&mut one.clone())
+        }
+        all
+    }
+
+    /*
+    #[test]
+    fn container() {
+        let data1 = hash(&vec![1; 32]);
+        let data2 = hash(&vec![2; 32]);
+        let data3 = hash(&vec![3; 32]);
+        let data4 = hash(&vec![4; 32]);
+
+        let data = join(vec![&data1, &data2, &data3, &data4]);
+
+        let cache = cache_builder(&data).unwrap();
+    }
+    */
+    #[test]
+    fn merkleize_4_leaves() {
+        let data1 = hash(&int_to_bytes32(1));
+        let data2 = hash(&int_to_bytes32(2));
+        let data3 = hash(&int_to_bytes32(3));
+        let data4 = hash(&int_to_bytes32(4));
+
+        let data = join(vec![&data1, &data2, &data3, &data4]);
+
+        let cache = merkleize(&data);
+
+        let hash_12 = {
+            let mut joined = vec![];
+            joined.append(&mut data1.clone());
+            joined.append(&mut data2.clone());
+            hash(&joined)
+        };
+        let hash_34 = {
+            let mut joined = vec![];
+            joined.append(&mut data3.clone());
+            joined.append(&mut data4.clone());
+            hash(&joined)
+        };
+        let hash_hash12_hash_34 = {
+            let mut joined = vec![];
+            joined.append(&mut hash_12.clone());
+            joined.append(&mut hash_34.clone());
+            hash(&joined)
+        };
+
+        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
+            let expected = match i {
+                0 => hash_hash12_hash_34.clone(),
+                1 => hash_12.clone(),
+                2 => hash_34.clone(),
+                3 => data1.clone(),
+                4 => data2.clone(),
+                5 => data3.clone(),
+                6 => data4.clone(),
+                _ => vec![],
+            };
+
+            assert_eq!(chunk, &expected[..], "failed at {}", i);
+        }
+    }
+}
+/*
+
 /// Get merkle root of some hashed values - the input leaf nodes is expected to already be hashed
 /// Outputs a `Vec<u8>` byte array of the merkle root given a set of leaf node values.
 pub fn cache_builder(values: &[u8]) -> Option<Vec<u8>> {
@@ -177,52 +303,51 @@ mod tests {
         let cache = cache_builder(&data).unwrap();
     }
     */
+#[test]
+fn can_build_cache() {
+let data1 = hash(&vec![1; 32]);
+let data2 = hash(&vec![2; 32]);
+let data3 = hash(&vec![3; 32]);
+let data4 = hash(&vec![4; 32]);
 
-    #[test]
+let data = join(vec![&data1, &data2, &data3, &data4]);
-    fn can_build_cache() {
-        let data1 = hash(&vec![1; 32]);
-        let data2 = hash(&vec![2; 32]);
-        let data3 = hash(&vec![3; 32]);
-        let data4 = hash(&vec![4; 32]);
 
-        let data = join(vec![&data1, &data2, &data3, &data4]);
+let cache = cache_builder(&data).unwrap();
 
-        let cache = cache_builder(&data).unwrap();
+let hash_12 = {
+let mut joined = vec![];
+joined.append(&mut data1.clone());
+joined.append(&mut data2.clone());
+hash(&joined)
+};
+let hash_34 = {
+let mut joined = vec![];
+joined.append(&mut data3.clone());
+joined.append(&mut data4.clone());
+hash(&joined)
+};
+let hash_hash12_hash_34 = {
+let mut joined = vec![];
+joined.append(&mut hash_12.clone());
+joined.append(&mut hash_34.clone());
+hash(&joined)
+};
 
-        let hash_12 = {
+for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
-            let mut joined = vec![];
+let expected = match i {
-            joined.append(&mut data1.clone());
+0 => hash_hash12_hash_34.clone(),
-            joined.append(&mut data2.clone());
+1 => hash_12.clone(),
-            hash(&joined)
+2 => hash_34.clone(),
-        };
+3 => data1.clone(),
-        let hash_34 = {
+4 => data2.clone(),
-            let mut joined = vec![];
+5 => data3.clone(),
-            joined.append(&mut data3.clone());
+6 => data4.clone(),
-            joined.append(&mut data4.clone());
+_ => vec![],
-            hash(&joined)
+};
-        };
-        let hash_hash12_hash_34 = {
-            let mut joined = vec![];
-            joined.append(&mut hash_12.clone());
-            joined.append(&mut hash_34.clone());
-            hash(&joined)
-        };
 
-        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
+assert_eq!(chunk, &expected[..], "failed at {}", i);
-            let expected = match i {
+}
-                0 => hash_hash12_hash_34.clone(),
+}
-                1 => hash_12.clone(),
-                2 => hash_34.clone(),
-                3 => data1.clone(),
-                4 => data2.clone(),
-                5 => data3.clone(),
-                6 => data4.clone(),
-                _ => vec![],
-            };
-
-            assert_eq!(chunk, &expected[..], "failed at {}", i);
-        }
-    }
 }
 
 /*
@@ -329,3 +454,4 @@ mod tests {
     }
 }
 */
+*/