From 5a2cf6219e129557c09e67842c74b03a17bf71cf Mon Sep 17 00:00:00 2001
From: Kirk Baird <baird.k@outlook.com>
Date: Fri, 25 Jan 2019 12:59:11 +1100
Subject: [PATCH] Modify tree_hash

---
 eth2/utils/ssz/src/impl_tree_hash.rs | 47 +---------------------------
 eth2/utils/ssz/src/tree_hash.rs      | 43 ++++++++++---------------
 2 files changed, 18 insertions(+), 72 deletions(-)

diff --git a/eth2/utils/ssz/src/impl_tree_hash.rs b/eth2/utils/ssz/src/impl_tree_hash.rs
index 7249b7f739..9463283cb5 100644
--- a/eth2/utils/ssz/src/impl_tree_hash.rs
+++ b/eth2/utils/ssz/src/impl_tree_hash.rs
@@ -1,6 +1,5 @@
 use super::ethereum_types::{Address, H256};
 use super::{hash, merkle_hash, ssz_encode, TreeHash};
-use std::collections::HashMap;
 
 impl TreeHash for u8 {
     fn hash_tree_root(&self) -> Vec<u8> {
@@ -67,57 +66,13 @@ where
     }
 }
 
-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)
-    }
-}
-
 #[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();
+        let result = vec![1u32, 2, 3, 4, 5, 6, 7].hash_tree_root();
         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
-            ]
-        );
-    }
-
 }
diff --git a/eth2/utils/ssz/src/tree_hash.rs b/eth2/utils/ssz/src/tree_hash.rs
index 6efe41a407..5f62c1c7b8 100644
--- a/eth2/utils/ssz/src/tree_hash.rs
+++ b/eth2/utils/ssz/src/tree_hash.rs
@@ -11,39 +11,31 @@ pub trait TreeHash {
 /// 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);
+    let (chunk_size, mut chunkz) = list_to_blob(list);
 
     // get data_len as bytes. It will hashed will the merkle root
     let datalen = list.len().to_le_bytes();
 
-    // merklize
-    let mut mhash = hash_level(&mut data, chunk_size);
-    while mhash.len() > HASHSIZE {
-        mhash = hash_level(&mut mhash, HASHSIZE);
-    }
+    // Tree-hash
+    while chunkz.len() > HASHSIZE {
+        let mut new_chunkz: Vec<u8> = Vec::new();
 
-    mhash.append(&mut datalen.to_vec());
-    hash(&mhash)
-}
-
-/// 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 hash(&c));
-        } else {
-            // Hash two chuncks together
-            result.append(&mut hash(two_chunks));
+        for two_chunks in chunkz.chunks(chunk_size * 2) {
+            if two_chunks.len() == chunk_size {
+                // Odd number of chunks
+                let mut c = two_chunks.to_vec();
+                c.append(&mut vec![0; SSZ_CHUNK_SIZE]);
+                new_chunkz.append(&mut hash(&c));
+            } else {
+                // Hash two chuncks together
+                new_chunkz.append(&mut hash(two_chunks));
+            }
         }
+        chunkz = new_chunkz;
     }
 
-    result
+    chunkz.append(&mut datalen.to_vec());
+    hash(&chunkz)
 }
 
 fn list_to_blob(list: &mut Vec<Vec<u8>>) -> (usize, Vec<u8>) {
@@ -68,7 +60,6 @@ fn list_to_blob(list: &mut Vec<Vec<u8>>) -> (usize, Vec<u8>) {
             data.append(item);
         }
     }
-
     (chunk_size, data)
 }