From 725401d2e1fc54e56fd60734d28e0612b062b0f8 Mon Sep 17 00:00:00 2001
From: Johns Beharry <johns@peakshift.com>
Date: Thu, 21 Mar 2019 02:51:49 +0100
Subject: [PATCH] test(hashing/merkle_root): write test for empty vector and
 odd leaf count

Signed-off-by: Johns Beharry <johns@peakshift.com>
---
 eth2/utils/hashing/src/lib.rs | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/eth2/utils/hashing/src/lib.rs b/eth2/utils/hashing/src/lib.rs
index 3f544d2c1..bbb37f6ab 100644
--- a/eth2/utils/hashing/src/lib.rs
+++ b/eth2/utils/hashing/src/lib.rs
@@ -8,11 +8,16 @@ pub fn hash(input: &[u8]) -> Vec<u8> {
     result
 }
 
-// Get merkle root of some hashed values - the input leaf nodes is expected to already be hashed
-// Outputs a `Vec<u8>` byte array of the merkle root given a set of leaf node values.
-pub fn merkle_root(values: &[Vec<u8>]) -> Vec<u8> {
+/// Get merkle root of some hashed values - the input leaf nodes is expected to already be hashed
+/// Outputs a `Vec<u8>` byte array of the merkle root given a set of leaf node values.
+pub fn merkle_root(values: &[Vec<u8>]) -> Option<Vec<u8>> {
     let values_len = values.len();
 
+    // check size of vector > 0 and ^ 2
+    if values.is_empty() || !values_len.is_power_of_two() {
+        return None
+    }
+
     // vector to store hashes
     // filled with 0 as placeholders
     let mut o: Vec<Vec<u8>> = vec![vec![0]; values_len];
@@ -30,7 +35,7 @@ pub fn merkle_root(values: &[Vec<u8>]) -> Vec<u8> {
     }
 
     // the root hash will be at index 1
-    o[1].clone()
+    return Some(o[1].clone())
 }
 
 #[cfg(test)]
@@ -79,6 +84,22 @@ mod tests {
 
         let expected = hash(&root[..]);
 
-        assert_eq!(&expected[..], output.as_slice());
+        assert_eq!(&expected[..], output.unwrap().as_slice());
+    }
+    #[test]
+    fn test_empty_input_merkle_root() {
+        let input = vec![];
+        let output = merkle_root(&input[..]);
+        assert_eq!(None, output);
+    }
+    #[test]
+    fn test_odd_leaf_merkle_root() {
+        let input = vec![
+            hash("a".as_bytes()),
+            hash("b".as_bytes()),
+            hash("a".as_bytes()),
+        ];
+        let output = merkle_root(&input[..]);
+        assert_eq!(None, output);
     }
 }