/
node.go
executable file
·201 lines (178 loc) · 5.06 KB
/
node.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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
package memiavl
import (
"bytes"
"crypto/sha256"
"encoding/binary"
"fmt"
"io"
)
// Node interface encapsulate the interface of both PersistedNode and MemNode.
type Node interface {
Height() uint8
IsLeaf() bool
Size() int64
Version() uint32
Key() []byte
Value() []byte
Left() Node
Right() Node
Hash() []byte
// SafeHash returns byte slice that's safe to retain
SafeHash() []byte
// PersistedNode clone a new node, MemNode modify in place
Mutate(version, cowVersion uint32) *MemNode
// Get query the value for a key, it's put into interface because a specialized implementation is more efficient.
Get(key []byte) ([]byte, uint32)
GetByIndex(uint32) ([]byte, []byte)
}
// setRecursive do set operation.
// it always do modification and return new `MemNode`, even if the value is the same.
// also returns if it's an update or insertion, if update, the tree height and balance is not changed.
func setRecursive(node Node, key, value []byte, version, cowVersion uint32) (*MemNode, bool) {
if node == nil {
return newLeafNode(key, value, version), true
}
nodeKey := node.Key()
if node.IsLeaf() {
switch bytes.Compare(key, nodeKey) {
case -1:
return &MemNode{
height: 1,
size: 2,
version: version,
key: nodeKey,
left: newLeafNode(key, value, version),
right: node,
}, false
case 1:
return &MemNode{
height: 1,
size: 2,
version: version,
key: key,
left: node,
right: newLeafNode(key, value, version),
}, false
default:
newNode := node.Mutate(version, cowVersion)
newNode.value = value
return newNode, true
}
} else {
var (
newChild, newNode *MemNode
updated bool
)
if bytes.Compare(key, nodeKey) == -1 {
newChild, updated = setRecursive(node.Left(), key, value, version, cowVersion)
newNode = node.Mutate(version, cowVersion)
newNode.left = newChild
} else {
newChild, updated = setRecursive(node.Right(), key, value, version, cowVersion)
newNode = node.Mutate(version, cowVersion)
newNode.right = newChild
}
if !updated {
newNode.updateHeightSize()
newNode = newNode.reBalance(version, cowVersion)
}
return newNode, updated
}
}
// removeRecursive returns:
// - (nil, origNode, nil) -> nothing changed in subtree
// - (value, nil, newKey) -> leaf node is removed
// - (value, new node, newKey) -> subtree changed
func removeRecursive(node Node, key []byte, version, cowVersion uint32) ([]byte, Node, []byte) {
if node == nil {
return nil, nil, nil
}
if node.IsLeaf() {
if bytes.Equal(node.Key(), key) {
return node.Value(), nil, nil
}
return nil, node, nil
}
if bytes.Compare(key, node.Key()) == -1 {
value, newLeft, newKey := removeRecursive(node.Left(), key, version, cowVersion)
if value == nil {
return nil, node, nil
}
if newLeft == nil {
return value, node.Right(), node.Key()
}
newNode := node.Mutate(version, cowVersion)
newNode.left = newLeft
newNode.updateHeightSize()
return value, newNode.reBalance(version, cowVersion), newKey
}
value, newRight, newKey := removeRecursive(node.Right(), key, version, cowVersion)
if value == nil {
return nil, node, nil
}
if newRight == nil {
return value, node.Left(), nil
}
newNode := node.Mutate(version, cowVersion)
newNode.right = newRight
if newKey != nil {
newNode.key = newKey
}
newNode.updateHeightSize()
return value, newNode.reBalance(version, cowVersion), nil
}
// Writes the node's hash to the given `io.Writer`. This function recursively calls
// children to update hashes.
func writeHashBytes(node Node, w io.Writer) error {
var (
n int
buf [binary.MaxVarintLen64]byte
)
n = binary.PutVarint(buf[:], int64(node.Height()))
if _, err := w.Write(buf[0:n]); err != nil {
return fmt.Errorf("writing height, %w", err)
}
n = binary.PutVarint(buf[:], node.Size())
if _, err := w.Write(buf[0:n]); err != nil {
return fmt.Errorf("writing size, %w", err)
}
n = binary.PutVarint(buf[:], int64(node.Version()))
if _, err := w.Write(buf[0:n]); err != nil {
return fmt.Errorf("writing version, %w", err)
}
// Key is not written for inner nodes, unlike writeBytes.
if node.IsLeaf() {
if err := EncodeBytes(w, node.Key()); err != nil {
return fmt.Errorf("writing key, %w", err)
}
// Indirection needed to provide proofs without values.
// (e.g. ProofLeafNode.ValueHash)
valueHash := sha256.Sum256(node.Value())
if err := EncodeBytes(w, valueHash[:]); err != nil {
return fmt.Errorf("writing value, %w", err)
}
} else {
if err := EncodeBytes(w, node.Left().Hash()); err != nil {
return fmt.Errorf("writing left hash, %w", err)
}
if err := EncodeBytes(w, node.Right().Hash()); err != nil {
return fmt.Errorf("writing right hash, %w", err)
}
}
return nil
}
// HashNode computes the hash of the node.
func HashNode(node Node) []byte {
if node == nil {
return nil
}
h := sha256.New()
if err := writeHashBytes(node, h); err != nil {
panic(err)
}
return h.Sum(nil)
}
// VerifyHash compare node's cached hash with computed one
func VerifyHash(node Node) bool {
return bytes.Equal(HashNode(node), node.Hash())
}