arbo/vt.go
vt.goRaw
1// Package arbo > vt.go implements the Virtual Tree, which computes a tree
2// without computing any hash. With the idea of once all the leafs are placed in
3// their positions, the hashes can be computed, avoiding computing a node hash
4// more than one time.
5package arbo
6
7import (
8 "bytes"
9 "encoding/hex"
10 "fmt"
11 "io"
12 "math"
13 "runtime"
14 "sync"
15)
16
17type node struct {
18 l *node
19 r *node
20 k []byte
21 v []byte
22 path []bool
23 h []byte
24}
25
26type params struct {
27 maxLevels int
28 hashFunction HashFunction
29 emptyHash []byte
30 dbg *dbgStats
31}
32
33func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, error) {
34 if len(ks) != len(vs) {
35 return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
36 len(ks), len(vs))
37 }
38 kvs := make([]kv, len(ks))
39 for i := 0; i < len(ks); i++ {
40 keyPath := make([]byte, p.hashFunction.Len())
41 copy(keyPath[:], ks[i])
42 kvs[i].pos = i
43 kvs[i].keyPath = keyPath
44 kvs[i].k = ks[i]
45 kvs[i].v = vs[i]
46 }
47
48 return kvs, nil
49}
50
51// vt stands for virtual tree. It's a tree that does not have any computed hash
52// while placing the leafs. Once all the leafs are placed, it computes all the
53// hashes. In this way, each node hash is only computed one time (at the end)
54// and the tree is computed in memory.
55type vt struct {
56 root *node
57 params *params
58}
59
60func newVT(maxLevels int, hash HashFunction) vt {
61 return vt{
62 root: nil,
63 params: ¶ms{
64 maxLevels: maxLevels,
65 hashFunction: hash,
66 emptyHash: make([]byte, hash.Len()), // empty
67 },
68 }
69}
70
71func (t *vt) addBatch(ks, vs [][]byte) error {
72 // parallelize adding leafs in the virtual tree
73 nCPU := flp2(runtime.NumCPU())
74 if nCPU == 1 || len(ks) < nCPU {
75 // var invalids []int
76 for i := 0; i < len(ks); i++ {
77 if err := t.add(0, ks[i], vs[i]); err != nil {
78 // invalids = append(invalids, i)
79 fmt.Println(err) // TODO WIP
80 }
81 }
82 return nil // TODO invalids
83 }
84
85 l := int(math.Log2(float64(nCPU)))
86
87 kvs, err := t.params.keysValuesToKvs(ks, vs)
88 if err != nil {
89 return err
90 }
91
92 buckets := splitInBuckets(kvs, nCPU)
93
94 nodesAtL, err := t.getNodesAtLevel(l)
95 if err != nil {
96 return err
97 }
98 // fmt.Println("nodesatL pre-E", len(nodesAtL))
99 if len(nodesAtL) != nCPU {
100 // CASE E: add one key at each bucket, and then do CASE D
101 for i := 0; i < len(buckets); i++ {
102 // add one leaf of the bucket, if there is an error when
103 // adding the k-v, try to add the next one of the bucket
104 // (until one is added)
105 var inserted int
106 for j := 0; j < len(buckets[i]); j++ {
107 if err := t.add(0, buckets[i][j].k, buckets[i][j].v); err == nil {
108 inserted = j
109 break
110 }
111 }
112
113 // remove the inserted element from buckets[i]
114 buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
115 }
116 nodesAtL, err = t.getNodesAtLevel(l)
117 if err != nil {
118 return err
119 }
120 }
121
122 subRoots := make([]*node, nCPU)
123 invalidsInBucket := make([][]int, nCPU)
124
125 var wg sync.WaitGroup
126 wg.Add(nCPU)
127 for i := 0; i < nCPU; i++ {
128 go func(cpu int) {
129 sortKvs(buckets[cpu])
130
131 bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
132 bucketVT.root = nodesAtL[cpu]
133 for j := 0; j < len(buckets[cpu]); j++ {
134 if err = bucketVT.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
135 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
136 }
137 }
138 subRoots[cpu] = bucketVT.root
139 wg.Done()
140 }(i)
141 }
142 wg.Wait()
143
144 newRootNode, err := upFromNodes(subRoots)
145 if err != nil {
146 return err
147 }
148 t.root = newRootNode
149
150 return nil
151}
152
153func (t *vt) getNodesAtLevel(l int) ([]*node, error) {
154 if t.root == nil {
155 return nil, nil
156 }
157 return t.root.getNodesAtLevel(0, l)
158}
159
160func (n *node) getNodesAtLevel(currLvl, l int) ([]*node, error) {
161 var nodes []*node
162
163 typ := n.typ()
164 if currLvl == l && typ != vtEmpty {
165 nodes = append(nodes, n)
166 return nodes, nil
167 }
168 if currLvl >= l {
169 panic("should not reach this point") // TODO TMP
170 // return nil, nil
171 }
172
173 if n.l != nil {
174 nodesL, err := n.l.getNodesAtLevel(currLvl+1, l)
175 if err != nil {
176 return nil, err
177 }
178 nodes = append(nodes, nodesL...)
179 }
180 if n.r != nil {
181 nodesR, err := n.r.getNodesAtLevel(currLvl+1, l)
182 if err != nil {
183 return nil, err
184 }
185 nodes = append(nodes, nodesR...)
186 }
187 return nodes, nil
188}
189
190func upFromNodes(ns []*node) (*node, error) {
191 if len(ns) == 1 {
192 return ns[0], nil
193 }
194
195 var res []*node
196 for i := 0; i < len(ns); i += 2 {
197 if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtEmpty {
198 // when both sub nodes are empty, the node is also empty
199 res = append(res, ns[i]) // empty node
200 }
201 n := &node{
202 l: ns[i],
203 r: ns[i+1],
204 }
205 res = append(res, n)
206 }
207 return upFromNodes(res)
208}
209
210func (t *vt) add(fromLvl int, k, v []byte) error {
211 leaf := newLeafNode(t.params, k, v)
212 if t.root == nil {
213 t.root = leaf
214 return nil
215 }
216
217 if err := t.root.add(t.params, fromLvl, leaf); err != nil {
218 return err
219 }
220
221 return nil
222}
223
224// computeHashes should be called after all the vt.add is used, once all the
225// leafs are in the tree
226func (t *vt) computeHashes() ([][2][]byte, error) {
227 var pairs [][2][]byte
228 var err error
229 // TODO parallelize computeHashes
230 pairs, err = t.root.computeHashes(t.params, pairs)
231 if err != nil {
232 return pairs, err
233 }
234 return pairs, nil
235}
236
237func newLeafNode(p *params, k, v []byte) *node {
238 keyPath := make([]byte, p.hashFunction.Len())
239 copy(keyPath[:], k)
240 path := getPath(p.maxLevels, keyPath)
241 n := &node{
242 k: k,
243 v: v,
244 path: path,
245 }
246 return n
247}
248
249type virtualNodeType int
250
251const (
252 vtEmpty = 0 // for convenience uses same value that PrefixValueEmpty
253 vtLeaf = 1 // for convenience uses same value that PrefixValueLeaf
254 vtMid = 2 // for convenience uses same value that PrefixValueIntermediate
255)
256
257func (n *node) typ() virtualNodeType {
258 if n == nil {
259 return vtEmpty // TODO decide if return 'vtEmpty' or an error
260 }
261 if n.l == nil && n.r == nil && n.k != nil {
262 return vtLeaf
263 }
264 if n.l != nil || n.r != nil {
265 return vtMid
266 }
267 return vtEmpty
268}
269
270func (n *node) add(p *params, currLvl int, leaf *node) error {
271 if currLvl > p.maxLevels-1 {
272 return fmt.Errorf("max virtual level %d", currLvl)
273 }
274
275 if n == nil {
276 // n = leaf // TMP!
277 return nil
278 }
279
280 t := n.typ()
281 switch t {
282 case vtMid:
283 if leaf.path[currLvl] {
284 //right
285 if n.r == nil {
286 // empty sub-node, add the leaf here
287 n.r = leaf
288 return nil
289 }
290 if err := n.r.add(p, currLvl+1, leaf); err != nil {
291 return err
292 }
293 } else {
294 if n.l == nil {
295 // empty sub-node, add the leaf here
296 n.l = leaf
297 return nil
298 }
299 if err := n.l.add(p, currLvl+1, leaf); err != nil {
300 return err
301 }
302 }
303 case vtLeaf:
304 if bytes.Equal(n.k, leaf.k) {
305 return fmt.Errorf("key already exists. Existing node: %s, trying to add node: %s",
306 hex.EncodeToString(n.k), hex.EncodeToString(leaf.k))
307 }
308
309 oldLeaf := &node{
310 k: n.k,
311 v: n.v,
312 path: n.path,
313 }
314 // remove values from current node (converting it to mid node)
315 n.k = nil
316 n.v = nil
317 n.h = nil
318 n.path = nil
319 if err := n.downUntilDivergence(p, currLvl, oldLeaf, leaf); err != nil {
320 return err
321 }
322 case vtEmpty:
323 panic(fmt.Errorf("EMPTY %v", n)) // TODO TMP
324 default:
325 return fmt.Errorf("ERR")
326 }
327
328 return nil
329}
330
331func (n *node) downUntilDivergence(p *params, currLvl int, oldLeaf, newLeaf *node) error {
332 if currLvl > p.maxLevels-1 {
333 return fmt.Errorf("max virtual level %d", currLvl)
334 }
335
336 if oldLeaf.path[currLvl] != newLeaf.path[currLvl] {
337 // reached divergence in next level
338 if newLeaf.path[currLvl] {
339 n.l = oldLeaf
340 n.r = newLeaf
341 } else {
342 n.l = newLeaf
343 n.r = oldLeaf
344 }
345 return nil
346 }
347 // no divergence yet, continue going down
348 if newLeaf.path[currLvl] {
349 // right
350 n.r = &node{}
351 if err := n.r.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
352 return err
353 }
354 } else {
355 // left
356 n.l = &node{}
357 if err := n.l.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
358 return err
359 }
360 }
361
362 return nil
363}
364
365// returns an array of key-values to store in the db
366func (n *node) computeHashes(p *params, pairs [][2][]byte) ([][2][]byte, error) {
367 if pairs == nil {
368 pairs = [][2][]byte{}
369 }
370 var err error
371 t := n.typ()
372 switch t {
373 case vtLeaf:
374 p.dbg.incHash()
375 leafKey, leafValue, err := newLeafValue(p.hashFunction, n.k, n.v)
376 if err != nil {
377 return pairs, err
378 }
379 n.h = leafKey
380 kv := [2][]byte{leafKey, leafValue}
381 pairs = append(pairs, kv)
382 case vtMid:
383 if n.l != nil {
384 pairs, err = n.l.computeHashes(p, pairs)
385 if err != nil {
386 return pairs, err
387 }
388 } else {
389 n.l = &node{
390 h: p.emptyHash,
391 }
392 }
393 if n.r != nil {
394 pairs, err = n.r.computeHashes(p, pairs)
395 if err != nil {
396 return pairs, err
397 }
398 } else {
399 n.r = &node{
400 h: p.emptyHash,
401 }
402 }
403 // once the sub nodes are computed, can compute the current node
404 // hash
405 p.dbg.incHash()
406 k, v, err := newIntermediate(p.hashFunction, n.l.h, n.r.h)
407 if err != nil {
408 return nil, err
409 }
410 n.h = k
411 kv := [2][]byte{k, v}
412 pairs = append(pairs, kv)
413 default:
414 return nil, fmt.Errorf("ERR TMP") // TODO
415 }
416
417 return pairs, nil
418}
419
420//nolint:unused
421func (t *vt) graphviz(w io.Writer) error {
422 fmt.Fprintf(w, `digraph hierarchy {
423node [fontname=Monospace,fontsize=10,shape=box]
424`)
425 if _, err := t.root.graphviz(w, t.params, 0); err != nil {
426 return err
427 }
428 fmt.Fprintf(w, "}\n")
429 return nil
430}
431
432//nolint:unused
433func (n *node) graphviz(w io.Writer, p *params, nEmpties int) (int, error) {
434 nChars := 4 // TODO move to global constant
435 if n == nil {
436 return nEmpties, nil
437 }
438
439 t := n.typ()
440 switch t {
441 case vtLeaf:
442 leafKey, _, err := newLeafValue(p.hashFunction, n.k, n.v)
443 if err != nil {
444 return nEmpties, err
445 }
446 fmt.Fprintf(w, "\"%p\" [style=filled,label=\"%v\"];\n", n, hex.EncodeToString(leafKey[:nChars]))
447
448 fmt.Fprintf(w, "\"%p\" -> {\"k:%v\\nv:%v\"}\n", n,
449 hex.EncodeToString(n.k[:nChars]),
450 hex.EncodeToString(n.v[:nChars]))
451 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
452 hex.EncodeToString(n.k[:nChars]),
453 hex.EncodeToString(n.v[:nChars]))
454 case vtMid:
455 fmt.Fprintf(w, "\"%p\" [label=\"\"];\n", n)
456
457 lStr := fmt.Sprintf("%p", n.l)
458 rStr := fmt.Sprintf("%p", n.r)
459 eStr := ""
460 if n.l == nil {
461 lStr = fmt.Sprintf("empty%v", nEmpties)
462 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
463 lStr)
464 nEmpties++
465 }
466 if n.r == nil {
467 rStr = fmt.Sprintf("empty%v", nEmpties)
468 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
469 rStr)
470 nEmpties++
471 }
472 fmt.Fprintf(w, "\"%p\" -> {\"%v\" \"%v\"}\n", n, lStr, rStr)
473 fmt.Fprint(w, eStr)
474 nEmpties, err := n.l.graphviz(w, p, nEmpties)
475 if err != nil {
476 return nEmpties, err
477 }
478 nEmpties, err = n.r.graphviz(w, p, nEmpties)
479 if err != nil {
480 return nEmpties, err
481 }
482
483 case vtEmpty:
484 default:
485 return nEmpties, fmt.Errorf("ERR")
486 }
487
488 return nEmpties, nil
489}
490
491//nolint:unused
492func (t *vt) printGraphviz() error {
493 w := bytes.NewBufferString("")
494 fmt.Fprintf(w,
495 "--------\nGraphviz:\n")
496 err := t.graphviz(w)
497 if err != nil {
498 fmt.Println(w)
499 return err
500 }
501 fmt.Fprintf(w,
502 "End of Graphviz --------\n")
503 fmt.Println(w)
504 return nil
505}