diff --git a/eth2/utils/ssz/Cargo.toml b/eth2/utils/ssz/Cargo.toml index f13db5def..21b726e9f 100644 --- a/eth2/utils/ssz/Cargo.toml +++ b/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" } diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index 7a5b1c527..99fd49221 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/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 { + self.chunk_modified.get(chunk).cloned() + } + + pub fn children_modified(&self, parent_chunk: usize) -> Option { + let children = children(parent_chunk); + + Some(self.changed(children.0)? | self.changed(children.1)?) + } + + pub fn hash_children(&self, parent_chunk: usize) -> Option> { + 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 { + 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( - &self, - other: &Self, - cache: &mut [u8], - i: usize, - changes: Vec, - ) -> Option<(usize, Vec)>; + fn build_cache_bytes(&self) -> Vec; + + fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>; } impl CachedTreeHash for u64 { - fn cached_hash_tree_root( - &self, - other: &Self, - cache: &mut [u8], - i: usize, - mut changes: Vec, - ) -> Option<(usize, Vec)> { - if self != other { - cache - .get_mut(i..i + HASHSIZE)? - .copy_from_slice(&mut hash(&ssz_encode(self))); - changes.push(true); - } else { - changes.push(false); - }; + fn build_cache_bytes(&self) -> Vec { + merkleize(&ssz_encode(self)) + } - 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( - &self, - other: &Self, - cache: &mut [u8], - i: usize, - mut changes: Vec, - ) -> Option<(usize, Vec)> { - let original_start = i; + fn build_cache_bytes(&self) -> Vec { + let mut leaves = vec![]; - let leaves = 4; - let nodes = num_nodes(leaves); - let internal = nodes - leaves; - let leaves_start = i + internal * HASHSIZE; + leaves.append(&mut self.a.build_cache_bytes()); + leaves.append(&mut self.b.build_cache_bytes()); + leaves.append(&mut self.c.build_cache_bytes()); + leaves.append(&mut self.d.build_cache_bytes()); - let mut leaf_changes = { - let leaf_changes = Vec::with_capacity(leaves); - let leaf_start = leaves_start; + merkleize(&leaves) + } - let (leaf_start, leaf_changes) = - self.a - .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)?; + fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> { + let num_leaves = 4; - 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); - changes.append(&mut leaf_changes); - - 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 + for chunk in 0..internal_chunks { + if cache.children_modified(chunk)? { + cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?; } } - Some((42, vec![any_changes])) + Some(()) } } +/// A reference function to test against. +pub fn merkleize(values: &[u8]) -> Vec { + let leaves = values.len() / HASHSIZE; + + if leaves == 0 || !leaves.is_power_of_two() { + panic!("Handle bad leaf count"); + } + + let mut o: Vec = 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 { + 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` byte array of the merkle root given a set of leaf node values. pub fn cache_builder(values: &[u8]) -> Option> { @@ -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] - 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 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 = { - 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![], +}; - 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); - } - } +assert_eq!(chunk, &expected[..], "failed at {}", i); +} +} } /* @@ -329,3 +454,4 @@ mod tests { } } */ +*/