Directory Restructure (#1163)

* Move tests -> testing

* Directory restructure

* Update Cargo.toml during restructure

* Update Makefile during restructure

* Fix arbitrary path

consensus/cached_tree_hash/src/cache.rs

@@ -0,0 +1,237 @@
+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;
+
+type CacheArena = cache_arena::CacheArena<Hash256>;
+type CacheArenaAllocation = cache_arena::CacheArenaAllocation<Hash256>;
+
+/// Sparse Merkle tree suitable for tree hashing vectors and lists.
+#[derive(Debug, PartialEq, Clone, Default, Encode, Decode)]
+pub struct TreeHashCache {
+    pub initialized: bool,
+    /// Depth is such that the tree has a capacity for 2^depth leaves
+    depth: usize,
+    /// Sparse layers.
+    ///
+    /// 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: 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 {
+        let mut layers = SmallVec8::with_capacity(depth + 1);
+
+        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,
+            depth,
+            layers,
+        }
+    }
+
+    /// Compute the updated Merkle root for the given `leaves`.
+    pub fn recalculate_merkle_root(
+        &mut self,
+        arena: &mut CacheArena,
+        leaves: impl Iterator<Item = [u8; BYTES_PER_CHUNK]> + ExactSizeIterator,
+    ) -> Result<Hash256, Error> {
+        let dirty_indices = self.update_leaves(arena, leaves)?;
+        self.update_merkle_root(arena, dirty_indices)
+    }
+
+    /// Phase 1 of the algorithm: compute the indices of all dirty leaves.
+    pub fn update_leaves(
+        &mut self,
+        arena: &mut CacheArena,
+        mut leaves: impl Iterator<Item = [u8; BYTES_PER_CHUNK]> + ExactSizeIterator,
+    ) -> Result<SmallVec8<usize>, Error> {
+        let new_leaf_count = leaves.len();
+
+        if new_leaf_count < self.leaves().len(arena)? {
+            return Err(Error::CannotShrink);
+        } else if new_leaf_count > 2usize.pow(self.depth as u32) {
+            return Err(Error::TooManyLeaves);
+        }
+
+        let mut dirty = SmallVec8::new();
+
+        // Update the existing leaves
+        self.leaves()
+            .iter_mut(arena)?
+            .enumerate()
+            .zip(&mut leaves)
+            .for_each(|((i, leaf), new_leaf)| {
+                if !self.initialized || leaf.as_bytes() != new_leaf {
+                    leaf.assign_from_slice(&new_leaf);
+                    dirty.push(i);
+                }
+            });
+
+        // Push the rest of the new leaves (if any)
+        dirty.extend(self.leaves().len(arena)?..new_leaf_count);
+        self.leaves()
+            .extend_with_vec(arena, leaves.map(|l| Hash256::from_slice(&l)).collect())?;
+
+        Ok(dirty)
+    }
+
+    /// Phase 2: propagate changes upwards from the leaves of the tree, and compute the root.
+    ///
+    /// Returns an error if `dirty_indices` is inconsistent with the cache.
+    pub fn update_merkle_root(
+        &mut self,
+        arena: &mut CacheArena,
+        mut dirty_indices: SmallVec8<usize>,
+    ) -> Result<Hash256, Error> {
+        if dirty_indices.is_empty() {
+            return Ok(self.root(arena));
+        }
+
+        let mut depth = self.depth;
+
+        while depth > 0 {
+            let new_dirty_indices = lift_dirty(&dirty_indices);
+
+            for &idx in &new_dirty_indices {
+                let left_idx = 2 * idx;
+                let right_idx = left_idx + 1;
+
+                let left = self.layers[depth]
+                    .get(arena, left_idx)?
+                    .ok_or_else(|| Error::MissingLeftIdx(left_idx))?;
+                let right = self.layers[depth]
+                    .get(arena, right_idx)?
+                    .copied()
+                    .unwrap_or_else(|| Hash256::from_slice(&ZERO_HASHES[self.depth - depth]));
+
+                let new_hash = hash32_concat(left.as_bytes(), right.as_bytes());
+
+                match self.layers[depth - 1].get_mut(arena, idx)? {
+                    Some(hash) => {
+                        hash.assign_from_slice(&new_hash);
+                    }
+                    None => {
+                        // Parent layer should already contain nodes for all non-dirty indices
+                        if idx != self.layers[depth - 1].len(arena)? {
+                            return Err(Error::CacheInconsistent);
+                        }
+                        self.layers[depth - 1].push(arena, Hash256::from_slice(&new_hash))?;
+                    }
+                }
+            }
+
+            dirty_indices = new_dirty_indices;
+            depth -= 1;
+        }
+
+        self.initialized = true;
+
+        Ok(self.root(arena))
+    }
+
+    /// Get the root of this cache, without doing any updates/computation.
+    pub fn root(&self, arena: &CacheArena) -> Hash256 {
+        self.layers[0]
+            .get(arena, 0)
+            .expect("cached tree should have a root layer")
+            .copied()
+            .unwrap_or_else(|| Hash256::from_slice(&ZERO_HASHES[self.depth]))
+    }
+
+    pub fn leaves(&mut self) -> &mut CacheArenaAllocation {
+        &mut self.layers[self.depth]
+    }
+}
+
+/// Compute the dirty indices for one layer up.
+fn lift_dirty(dirty_indices: &[usize]) -> SmallVec8<usize> {
+    let mut new_dirty = SmallVec8::with_capacity(dirty_indices.len());
+
+    for index in dirty_indices {
+        new_dirty.push(index / 2)
+    }
+
+    new_dirty.dedup();
+    new_dirty
+}
+
+/// Returns the number of nodes that should be at each layer of a tree with the given `depth` and
+/// number of `leaves`.
+///
+/// Note: the top-most layer is `0` and a tree that has 8 leaves (4 layers) has a depth of 3 (_not_
+/// a depth of 4).
+///
+/// ## Example
+///
+/// Consider the following tree that has `depth = 3` and `leaves = 5`.
+///
+///```ignore
+/// 0        o      <-- height 0 has 1 node
+///        /   \
+/// 1    o      o   <-- height 1 has 2 nodes
+///     / \    /
+/// 2  o   o   o    <-- height 2 has 3 nodes
+///   /\   /\ /
+/// 3 o o o o o     <-- height 3 have 5 nodes
+/// ```
+fn nodes_per_layer(layer: usize, depth: usize, leaves: usize) -> usize {
+    if layer == depth {
+        leaves
+    } else {
+        let leaves_per_node = 1 << (depth - layer);
+        (leaves + leaves_per_node - 1) / leaves_per_node
+    }
+}
+
+#[cfg(test)]
+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);
+        assert_eq!(nodes_per_layer(1, 3, 5), 2);
+        assert_eq!(nodes_per_layer(2, 3, 5), 3);
+        assert_eq!(nodes_per_layer(3, 3, 5), 5);
+    }
+
+    #[test]
+    fn test_node_per_layer_balanced_tree() {
+        assert_eq!(nodes_per_layer(0, 3, 8), 1);
+        assert_eq!(nodes_per_layer(1, 3, 8), 2);
+        assert_eq!(nodes_per_layer(2, 3, 8), 4);
+        assert_eq!(nodes_per_layer(3, 3, 8), 8);
+    }
+}