From c961c87307b89cdff552a1437b3699bf09f91aea Mon Sep 17 00:00:00 2001
From: mjkeating <michael.j.keating@gmail.com>
Date: Sun, 9 Dec 2018 10:12:41 -0800
Subject: [PATCH] fixed tree_hash() for Vec<T> and a couple of other issues

---
 beacon_chain/types/src/shard_and_committee.rs |  2 +-
 beacon_chain/types/src/validator_record.rs    |  2 +-
 beacon_chain/utils/ssz/src/impl_tree_hash.rs  | 16 +++++++---------
 beacon_chain/utils/ssz/src/tree_hash.rs       | 19 ++++++-------------
 4 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/beacon_chain/types/src/shard_and_committee.rs b/beacon_chain/types/src/shard_and_committee.rs
index 9982611e8..c7fed2e8d 100644
--- a/beacon_chain/types/src/shard_and_committee.rs
+++ b/beacon_chain/types/src/shard_and_committee.rs
@@ -29,7 +29,7 @@ impl TreeHash for ShardAndCommittee {
         result.append(&mut self.shard.tree_hash());
         result.append(&mut merkle_hash(&mut committee_ssz_items));
 
-        result.tree_hash()
+        result.as_slice().tree_hash()
     }
 }
 
diff --git a/beacon_chain/types/src/validator_record.rs b/beacon_chain/types/src/validator_record.rs
index 1b18a0643..799e0137f 100644
--- a/beacon_chain/types/src/validator_record.rs
+++ b/beacon_chain/types/src/validator_record.rs
@@ -77,7 +77,7 @@ impl TreeHash for ValidatorRecord {
         balance.resize(16, 0);
         ssz.append(&mut balance);
 
-        ssz.tree_hash()
+        ssz.as_slice().tree_hash()
     }
 }
 
diff --git a/beacon_chain/utils/ssz/src/impl_tree_hash.rs b/beacon_chain/utils/ssz/src/impl_tree_hash.rs
index 8c5a5d17a..dfa6da2b2 100644
--- a/beacon_chain/utils/ssz/src/impl_tree_hash.rs
+++ b/beacon_chain/utils/ssz/src/impl_tree_hash.rs
@@ -2,7 +2,7 @@ extern crate blake2_rfc;
 
 use self::blake2_rfc::blake2b::blake2b;
 use super::ethereum_types::{Address, H256};
-use super::{ssz_encode, TreeHash};
+use super::{merkle_hash, ssz_encode, TreeHash};
 use std::cmp::Ord;
 use std::collections::HashMap;
 use std::hash::Hash;
@@ -53,15 +53,13 @@ impl<T> TreeHash for Vec<T>
 where
     T: TreeHash,
 {
-    /// Appends the tree_hash for each value of 'self' into a byte array
-    /// and returns the hash of said byte array
+    /// 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 result = Vec::new();
-        for x in self {
-            result.append(&mut x.tree_hash());
-        }
-
-        hash(&result)
+        let mut tree_hashes = self.iter().map(|x| x.tree_hash()).collect();
+        merkle_hash(&mut tree_hashes)
     }
 }
 
diff --git a/beacon_chain/utils/ssz/src/tree_hash.rs b/beacon_chain/utils/ssz/src/tree_hash.rs
index 60a052b35..31e1d3720 100644
--- a/beacon_chain/utils/ssz/src/tree_hash.rs
+++ b/beacon_chain/utils/ssz/src/tree_hash.rs
@@ -23,14 +23,12 @@ pub fn merkle_hash(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
     }
 
     mhash.append(data_len_bytes);
-    mhash.tree_hash()
+    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> {
-    assert!(data.len() % chunk_size == 0);
-
     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 {
@@ -38,7 +36,7 @@ fn hash_level(data: &mut Vec<u8>, chunk_size: usize) -> Vec<u8> {
             // CHUNKSIZE vector
             let mut c = two_chunks.to_vec();
             c.append(&mut vec![0; CHUNKSIZE]);
-            result.append(&mut c.tree_hash());
+            result.append(&mut c.as_slice().tree_hash());
         } else {
             result.append(&mut two_chunks.tree_hash());
         }
@@ -48,14 +46,14 @@ fn hash_level(data: &mut Vec<u8>, chunk_size: usize) -> Vec<u8> {
 }
 
 fn list_to_blob(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
-    let mut data_len = 0;
     if list[0].len().is_power_of_two() == false {
         for x in list.iter_mut() {
             extend_to_power_of_2(x);
-            data_len += x.len();
         }
     }
 
+    let mut data_len = list[0].len() * list.len();
+
     // do we need padding?
     let extend_by = if data_len % CHUNKSIZE > 0 {
         CHUNKSIZE - (data_len % CHUNKSIZE)
@@ -63,6 +61,8 @@ fn list_to_blob(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
         0
     };
 
+    println!("data_len {}, extend_by {}", data_len, extend_by);
+
     // allocate buffer and append each list element (flatten the vec of vecs)
     data_len += extend_by;
     let mut data: Vec<u8> = Vec::with_capacity(data_len);
@@ -93,13 +93,6 @@ fn extend_to_power_of_2(data: &mut Vec<u8>) {
 mod tests {
     use super::*;
 
-    #[test]
-    fn test_extend_to_power_of_2() {
-        let mut data = vec![1, 2, 3, 4, 5];
-        extend_to_power_of_2(&mut data);
-        assert_eq!(data, [1, 2, 3, 4, 5, 0, 0, 0]);
-    }
-
     #[test]
     fn test_merkle_hash() {
         let data1 = vec![1; 100];