From 5e43c770eefec45373382b76e939788a0c493ba5 Mon Sep 17 00:00:00 2001
From: Wei Yang <richard.weiyang@gmail.com>
Date: Mon, 16 Oct 2017 12:41:40 +0800
Subject: [PATCH] merkle tree: fix the implementation and add a test case

Merkle tree requires a full binary tree, which means the number of node is
power of 2 and the level of the tree is log(N) instead of N/2.

This patch duplicate the node to a power of 2 number and correct the level
calculation. Besides it adds a test case with 6 nodes.

Signed-off-by: Wei Yang <richard.weiyang@gmail.com>
---
 merkle_tree.go      |  6 ++++--
 merkle_tree_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/merkle_tree.go b/merkle_tree.go
index 7a4156b..0b7deb6 100644
--- a/merkle_tree.go
+++ b/merkle_tree.go
@@ -20,7 +20,8 @@ type MerkleNode struct {
 func NewMerkleTree(data [][]byte) *MerkleTree {
 	var nodes []MerkleNode
 
-	if len(data)%2 != 0 {
+	// append node until number is power of 2
+	for (len(data) & (len(data) - 1)) != 0 {
 		data = append(data, data[len(data)-1])
 	}
 
@@ -29,7 +30,8 @@ func NewMerkleTree(data [][]byte) *MerkleTree {
 		nodes = append(nodes, *node)
 	}
 
-	for i := 0; i < len(data)/2; i++ {
+	// up level until there is only 1 node
+	for len(nodes) > 1 {
 		var newLevel []MerkleNode
 
 		for j := 0; j < len(nodes); j += 2 {
diff --git a/merkle_tree_test.go b/merkle_tree_test.go
index acff5ff..d104cee 100644
--- a/merkle_tree_test.go
+++ b/merkle_tree_test.go
@@ -73,3 +73,41 @@ func TestNewMerkleTree(t *testing.T) {
 
 	assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct")
 }
+
+func TestNewMerkleTree8(t *testing.T) {
+	data := [][]byte{
+		[]byte("node1"),
+		[]byte("node2"),
+		[]byte("node3"),
+		[]byte("node4"),
+		[]byte("node5"),
+		[]byte("node6"),
+	}
+	// Level 0
+	n01 := NewMerkleNode(nil, nil, data[0])
+	n02 := NewMerkleNode(nil, nil, data[1])
+	n03 := NewMerkleNode(nil, nil, data[2])
+	n04 := NewMerkleNode(nil, nil, data[3])
+	n05 := NewMerkleNode(nil, nil, data[4])
+	n06 := NewMerkleNode(nil, nil, data[5])
+	n07 := NewMerkleNode(nil, nil, data[5])
+	n08 := NewMerkleNode(nil, nil, data[5])
+
+	// Level 1
+	n11 := NewMerkleNode(n01, n02, nil)
+	n12 := NewMerkleNode(n03, n04, nil)
+	n13 := NewMerkleNode(n05, n06, nil)
+	n14 := NewMerkleNode(n07, n08, nil)
+
+	// Level 2
+	n21 := NewMerkleNode(n11, n12, nil)
+	n22 := NewMerkleNode(n13, n14, nil)
+
+	// Level 3
+	n31 := NewMerkleNode(n21, n22, nil)
+
+	rootHash := fmt.Sprintf("%x", n31.Data)
+	mTree := NewMerkleTree(data)
+
+	assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct")
+}