diff --git a/merkle_tree.go b/merkle_tree.go index 7a4156bd..5f06f186 100644 --- a/merkle_tree.go +++ b/merkle_tree.go @@ -20,20 +20,20 @@ type MerkleNode struct { func NewMerkleTree(data [][]byte) *MerkleTree { var nodes []MerkleNode - if len(data)%2 != 0 { - data = append(data, data[len(data)-1]) - } - for _, datum := range data { node := NewMerkleNode(nil, nil, datum) nodes = append(nodes, *node) } - for i := 0; i < len(data)/2; i++ { + for len(nodes) > 1 { var newLevel []MerkleNode - for j := 0; j < len(nodes); j += 2 { - node := NewMerkleNode(&nodes[j], &nodes[j+1], nil) + for i := 0; i < len(nodes); i += 2 { + j := i + 1 + if j == len(nodes) { + j -= 1 + } + node := NewMerkleNode(&nodes[i], &nodes[j], nil) newLevel = append(newLevel, *node) } diff --git a/merkle_tree_test.go b/merkle_tree_test.go index acff5ffa..c3c9ad8b 100644 --- a/merkle_tree_test.go +++ b/merkle_tree_test.go @@ -73,3 +73,36 @@ func TestNewMerkleTree(t *testing.T) { assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct") } + +func TestNewMerkleTree2(t *testing.T) { + data := [][]byte{ + []byte("node1"), + []byte("node2"), + []byte("node3"), + []byte("node4"), + []byte("node5"), + } + // Level 1 + n11 := NewMerkleNode(nil, nil, data[0]) + n12 := NewMerkleNode(nil, nil, data[1]) + n13 := NewMerkleNode(nil, nil, data[2]) + n14 := NewMerkleNode(nil, nil, data[3]) + n15 := NewMerkleNode(nil, nil, data[4]) + + // Level 2 + n21 := NewMerkleNode(n11, n12, nil) + n22 := NewMerkleNode(n13, n14, nil) + n23 := NewMerkleNode(n15, n15, nil) + + // Level 3 + n31 := NewMerkleNode(n21, n22, nil) + n32 := NewMerkleNode(n23, n23, nil) + + // Level 4 + n41 := NewMerkleNode(n31, n32, nil) + + rootHash := fmt.Sprintf("%x", n41.Data) + mTree := NewMerkleTree(data) + + assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct") +}