-
Notifications
You must be signed in to change notification settings - Fork 0
/
tree.go
94 lines (77 loc) · 1.8 KB
/
tree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
package merkle
// Tree represents a merkle tree
type Tree struct {
// levels contains all the levels with the leaf nodes at the top i.e. 0 index
// and the root node at the end of the slice
levels [][]*Node
}
// GenerateTree generates a merkle tree from the given data.
func GenerateTree(data [][]byte) *Tree {
xtras := 0
leafs, added := generateLeafNodes(data)
if added {
xtras++
}
tree := &Tree{levels: [][]*Node{leafs}}
if len(data) == 1 {
return tree
}
for {
ll := tree.levels[len(tree.levels)-1]
pnodes, added := generateParentLevel(ll)
if added {
xtras++
}
tree.levels = append(tree.levels, pnodes)
if len(pnodes) == 1 {
break
}
}
return tree
}
// Root returns the root node of the tree
func (t *Tree) Root() *Node {
rl := t.levels[len(t.levels)-1]
return rl[0]
}
// Height returns the height of the tree
func (t *Tree) Height() int {
return len(t.levels)
}
// Leafs returns theh leaf nodes
func (t *Tree) Leafs() []*Node {
return t.levels[0]
}
// If data is odd, then the last item is repeated
func generateLeafNodes(data [][]byte) ([]*Node, bool) {
l := len(data)
lnodes := make([]*Node, l)
for i, d := range data {
lnodes[i] = NewNode(d)
}
if (l % 2) != 0 {
//return append(lnodes, NewNode(nil))
return append(lnodes, lnodes[l-1]), true
}
return lnodes, false
}
// If ns is odd the last node is repeated.
func generateParentLevel(ns []*Node) ([]*Node, bool) {
nodes := ns
added := false
if (len(ns) % 2) != 0 {
nodes = append(nodes, nodes[len(nodes)-1])
added = true
//nodes = append(nodes, NewNode(nil))
}
ln := len(nodes)
parent := make([]*Node, ln/2)
for i := 0; i < ln; i += 2 {
data := append(nodes[i].Hash(), nodes[i+1].Hash()...)
nn := NewNode(data)
nn.Left = nodes[i]
nn.Right = nodes[i+1]
parent[i/2] = nn
}
return parent, added
}