From da81f82e130b97fddf88f12b80486bf1442d6417 Mon Sep 17 00:00:00 2001
From: Vincent Batts <vbatts@hashbangbash.com>
Date: Tue, 17 Mar 2015 17:27:39 -0400
Subject: [PATCH] .: checksum for an unbalanced tree

---
 node_test.go | 48 ++++++++++++------------------------------------
 tree.go      | 28 +++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 37 deletions(-)

diff --git a/node_test.go b/node_test.go
index 6655c22..0ed24f1 100644
--- a/node_test.go
+++ b/node_test.go
@@ -1,7 +1,7 @@
 package merkle
 
 import (
-	"log"
+	"fmt"
 	"strings"
 	"testing"
 )
@@ -10,8 +10,9 @@ var words string = `Who were expelled from the academy for crazy & publishing ob
 
 func TestNodeSums(t *testing.T) {
 	var (
-		nodes []*Node
-		h     = DefaultHash.New()
+		nodes            []*Node
+		h                = DefaultHash.New()
+		expectedChecksum = "819fe8fed7a46900bd0613344c5ba2be336c74db"
 	)
 	for _, word := range strings.Split(words, " ") {
 		h.Reset()
@@ -28,39 +29,14 @@ func TestNodeSums(t *testing.T) {
 			break
 		}
 	}
-	for i := range nodes {
-		c, err := nodes[i].Checksum()
-		if err != nil {
-			t.Error(err)
-		}
-		t.Logf("checksum %x", c)
+	if len(nodes) != 1 {
+		t.Errorf("%d nodes", len(nodes))
 	}
-	if len(nodes) > 0 {
-		t.Errorf("%d nodes; %d characters", len(nodes), len(words))
+	c, err := nodes[0].Checksum()
+	if err != nil {
+		t.Error(err)
+	}
+	if gotChecksum := fmt.Sprintf("%x", c); gotChecksum != expectedChecksum {
+		t.Errorf("expected checksum %q, got %q", expectedChecksum, gotChecksum)
 	}
 }
-
-func levelUp(nodes []*Node) []*Node {
-	var (
-		newNodes []*Node
-		last     = len(nodes) - 1
-	)
-
-	for i := range nodes {
-		if i%2 == 0 {
-			if i == last {
-				// TODO rebalance the last parent
-				log.Println("WHOOP")
-			}
-			n := NewNode()
-			n.Left = nodes[i]
-			n.Left.Parent = n
-			newNodes = append(newNodes, n)
-		} else {
-			n := newNodes[len(newNodes)-1]
-			n.Right = nodes[i]
-			n.Right.Parent = n
-		}
-	}
-	return newNodes
-}
diff --git a/tree.go b/tree.go
index e65b2c5..bbd5683 100644
--- a/tree.go
+++ b/tree.go
@@ -1,5 +1,31 @@
 package merkle
 
 type Tree struct {
-	Nodes []Node
+	Nodes []*Node
+}
+
+func levelUp(nodes []*Node) []*Node {
+	var (
+		newNodes []*Node
+		last     = len(nodes) - 1
+	)
+
+	for i := range nodes {
+		if i%2 == 0 {
+			if i == last {
+				// last nodes on uneven node counts get pushed up, to be in the next level up
+				newNodes = append(newNodes, nodes[i])
+				continue
+			}
+			n := NewNode()
+			n.Left = nodes[i]
+			n.Left.Parent = n
+			newNodes = append(newNodes, n)
+		} else {
+			n := newNodes[len(newNodes)-1]
+			n.Right = nodes[i]
+			n.Right.Parent = n
+		}
+	}
+	return newNodes
 }