Merge pull request #169 from mjkeating/tree_hash2

Added TreeHash functionality

eth2/utils/ssz/Cargo.toml

@@ -7,3 +7,4 @@ edition = "2018"
 [dependencies]
 bytes = "0.4.9"
 ethereum-types = "0.4.0"
+hashing = { path = "../hashing" }

eth2/utils/ssz/src/impl_tree_hash.rs

@@ -0,0 +1,123 @@
+extern crate hashing;
+
+use self::hashing::canonical_hash;
+use super::ethereum_types::{Address, H256};
+use super::{merkle_hash, ssz_encode, TreeHash};
+use std::cmp::Ord;
+use std::collections::HashMap;
+use std::hash::Hash;
+
+impl TreeHash for u8 {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for u16 {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for u32 {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for u64 {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for Address {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for H256 {
+    fn tree_hash(&self) -> Vec<u8> {
+        ssz_encode(self)
+    }
+}
+
+impl TreeHash for [u8] {
+    fn tree_hash(&self) -> Vec<u8> {
+        hash(&self)
+    }
+}
+
+impl<T> TreeHash for Vec<T>
+where
+    T: TreeHash,
+{
+    /// Returns the merkle_hash of a list of tree_hash values created
+    /// from the given list.
+    /// Note: A byte vector, Vec<u8>, must be converted to a slice (as_slice())
+    ///       to be handled properly (i.e. hashed) as byte array.
+    fn tree_hash(&self) -> Vec<u8> {
+        let mut tree_hashes = self.iter().map(|x| x.tree_hash()).collect();
+        merkle_hash(&mut tree_hashes)
+    }
+}
+
+impl<K, V> TreeHash for HashMap<K, V>
+where
+    K: Eq,
+    K: Hash,
+    K: Ord,
+    V: TreeHash,
+{
+    /// Appends the tree_hash for each value of 'self, sorted by key,
+    /// into a byte array and returns the hash of said byte array
+    fn tree_hash(&self) -> Vec<u8> {
+        let mut items: Vec<_> = self.iter().collect();
+        items.sort_by(|a, b| a.0.cmp(b.0));
+        let mut result = Vec::new();
+        for item in items {
+            result.append(&mut item.1.tree_hash());
+        }
+
+        hash(&result)
+    }
+}
+
+fn hash(data: &[u8]) -> Vec<u8> {
+    canonical_hash(data)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_impl_tree_hash_vec() {
+        let result = vec![1u32, 2, 3, 4, 5, 6, 7].tree_hash();
+        assert_eq!(result.len(), 32);
+    }
+
+    #[test]
+    fn test_impl_tree_hash_hashmap() {
+        let mut map = HashMap::new();
+        map.insert("c", 3);
+        map.insert("b", 2);
+        map.insert("g", 7);
+        map.insert("d", 6);
+        map.insert("e", 4);
+        map.insert("a", 1u32);
+        map.insert("f", 5);
+        let result = map.tree_hash();
+
+        // TODO: create tests that tie-out to an offical result
+        assert_eq!(
+            result,
+            [
+                130, 215, 165, 255, 224, 6, 144, 225, 14, 139, 67, 238, 205, 240, 20, 173, 53, 0,
+                105, 62, 49, 174, 244, 160, 114, 92, 232, 11, 102, 200, 112, 24
+            ]
+        );
+    }
+
+}

eth2/utils/ssz/src/lib.rs

@@ -12,12 +12,15 @@ extern crate ethereum_types;
 
 pub mod decode;
 pub mod encode;
+pub mod tree_hash;
 
 mod impl_decode;
 mod impl_encode;
+mod impl_tree_hash;
 
 pub use crate::decode::{decode_ssz, decode_ssz_list, Decodable, DecodeError};
 pub use crate::encode::{Encodable, SszStream};
+pub use crate::tree_hash::{merkle_hash, TreeHash};
 
 pub const LENGTH_BYTES: usize = 4;
 pub const MAX_LIST_SIZE: usize = 1 << (4 * 8);

eth2/utils/ssz/src/tree_hash.rs

@@ -0,0 +1,89 @@
+const SSZ_CHUNK_SIZE: usize = 128;
+const HASHSIZE: usize = 32;
+
+pub trait TreeHash {
+    fn tree_hash(&self) -> Vec<u8>;
+}
+
+/// Returns a 32 byte hash of 'list' - a vector of byte vectors.
+/// Note that this will consume 'list'.
+pub fn merkle_hash(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
+    // flatten list
+    let (chunk_size, mut data) = list_to_blob(list);
+
+    // get data_len as bytes. It will hashed will the merkle root
+    let dlen = list.len() as u64;
+    let data_len_bytes = &mut dlen.tree_hash();
+    data_len_bytes.resize(32, 0);
+
+    // merklize
+    let mut mhash = hash_level(&mut data, chunk_size);
+    while mhash.len() > HASHSIZE {
+        mhash = hash_level(&mut mhash, HASHSIZE);
+    }
+
+    mhash.append(data_len_bytes);
+    mhash.as_slice().tree_hash()
+}
+
+/// Takes a flat vector of bytes. It then hashes 'chunk_size * 2' slices into
+/// a byte vector of hashes, divisible by HASHSIZE
+fn hash_level(data: &mut Vec<u8>, chunk_size: usize) -> Vec<u8> {
+    let mut result: Vec<u8> = Vec::new();
+    for two_chunks in data.chunks(chunk_size * 2) {
+        if two_chunks.len() == chunk_size && data.len() > chunk_size {
+            // if there is only one chunk here, hash it with a zero-byte
+            // SSZ_CHUNK_SIZE vector
+            let mut c = two_chunks.to_vec();
+            c.append(&mut vec![0; SSZ_CHUNK_SIZE]);
+            result.append(&mut c.as_slice().tree_hash());
+        } else {
+            result.append(&mut two_chunks.tree_hash());
+        }
+    }
+
+    result
+}
+
+fn list_to_blob(list: &mut Vec<Vec<u8>>) -> (usize, Vec<u8>) {
+    let chunk_size = if list.is_empty() {
+        SSZ_CHUNK_SIZE
+    } else if list[0].len() < SSZ_CHUNK_SIZE {
+        let items_per_chunk = SSZ_CHUNK_SIZE / list[0].len();
+        items_per_chunk * list[0].len()
+    } else {
+        list[0].len()
+    };
+
+    let mut data = Vec::new();
+    if list.is_empty() {
+        // handle and empty list
+        data.append(&mut vec![0; SSZ_CHUNK_SIZE]);
+    } else {
+        // just create a blob here; we'll divide into
+        // chunked slices when we merklize
+        data.reserve(list[0].len() * list.len());
+        for item in list.iter_mut() {
+            data.append(item);
+        }
+    }
+
+    (chunk_size, data)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_merkle_hash() {
+        let data1 = vec![1; 100];
+        let data2 = vec![2; 100];
+        let data3 = vec![3; 100];
+        let mut list = vec![data1, data2, data3];
+        let result = merkle_hash(&mut list);
+
+        //note: should test againt a known test hash value
+        assert_eq!(HASHSIZE, result.len());
+    }
+}