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) ([]int, 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 }
80 }
81 return invalids, nil
82 }
83
84 l := int(math.Log2(float64(nCPU)))
85
86 kvs, err := t.params.keysValuesToKvs(ks, vs)
87 if err != nil {
88 return nil, err
89 }
90
91 buckets := splitInBuckets(kvs, nCPU)
92
93 nodesAtL, err := t.getNodesAtLevel(l)
94 if err != nil {
95 return nil, err
96 }
97 if len(nodesAtL) != nCPU && t.root != nil {
98 // CASE E: add one key at each bucket, and then do CASE D
99 for i := 0; i < len(buckets); i++ {
100 // add one leaf of the bucket, if there is an error when
101 // adding the k-v, try to add the next one of the bucket
102 // (until one is added)
103 inserted := -1
104 for j := 0; j < len(buckets[i]); j++ {
105 if err := t.add(0, buckets[i][j].k, buckets[i][j].v); err == nil {
106 inserted = j
107 break
108 }
109 }
110
111 // remove the inserted element from buckets[i]
112 // fmt.Println("rm-ins", inserted)
113 if inserted != -1 {
114 buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
115 }
116 }
117 nodesAtL, err = t.getNodesAtLevel(l)
118 if err != nil {
119 return nil, err
120 }
121 }
122 if len(nodesAtL) != nCPU {
123 fmt.Println("ASDF")
124 panic("should not happen")
125 }
126
127 subRoots := make([]*node, nCPU)
128 invalidsInBucket := make([][]int, nCPU)
129
130 var wg sync.WaitGroup
131 wg.Add(nCPU)
132 for i := 0; i < nCPU; i++ {
133 go func(cpu int) {
134 sortKvs(buckets[cpu])
135
136 bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
137 bucketVT.root = nodesAtL[cpu]
138 for j := 0; j < len(buckets[cpu]); j++ {
139 if err = bucketVT.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
140 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
141 }
142 }
143 subRoots[cpu] = bucketVT.root
144 wg.Done()
145 }(i)
146 }
147 wg.Wait()
148
149 var invalids []int
150 for i := 0; i < len(invalidsInBucket); i++ {
151 invalids = append(invalids, invalidsInBucket[i]...)
152 }
153
154 newRootNode, err := upFromNodes(subRoots)
155 if err != nil {
156 return nil, err
157 }
158 t.root = newRootNode
159
160 return invalids, nil
161}
162
163func (t *vt) getNodesAtLevel(l int) ([]*node, error) {
164 if t.root == nil {
165 var r []*node
166 nChilds := int(math.Pow(2, float64(l))) //nolint:gomnd
167 for i := 0; i < nChilds; i++ {
168 r = append(r, nil)
169 }
170 return r, nil
171 }
172 return t.root.getNodesAtLevel(0, l)
173}
174
175func (n *node) getNodesAtLevel(currLvl, l int) ([]*node, error) {
176 if n == nil {
177 var r []*node
178 nChilds := int(math.Pow(2, float64(l-currLvl))) //nolint:gomnd
179 for i := 0; i < nChilds; i++ {
180 r = append(r, nil)
181 }
182 return r, nil
183 }
184
185 typ := n.typ()
186 if currLvl == l && typ != vtEmpty {
187 return []*node{n}, nil
188 }
189 if currLvl >= l {
190 panic("should not reach this point") // TODO TMP
191 }
192
193 var nodes []*node
194
195 nodesL, err := n.l.getNodesAtLevel(currLvl+1, l)
196 if err != nil {
197 return nil, err
198 }
199 nodes = append(nodes, nodesL...)
200
201 nodesR, err := n.r.getNodesAtLevel(currLvl+1, l)
202 if err != nil {
203 return nil, err
204 }
205 nodes = append(nodes, nodesR...)
206
207 return nodes, nil
208}
209
210func upFromNodes(ns []*node) (*node, error) {
211 if len(ns) == 1 {
212 return ns[0], nil
213 }
214
215 var res []*node
216 for i := 0; i < len(ns); i += 2 {
217 // if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtEmpty {
218 if ns[i] == nil && ns[i+1] == nil {
219 // when both sub nodes are empty, the node is also empty
220 res = append(res, ns[i]) // empty node
221 continue
222 }
223 n := &node{
224 l: ns[i],
225 r: ns[i+1],
226 }
227 res = append(res, n)
228 }
229 return upFromNodes(res)
230}
231
232// func upFromNodesComputingHashes(p *params, ns []*node, pairs [][2][]byte) (
233// [][2][]byte, *node, error) {
234// if len(ns) == 1 {
235// return pairs, ns[0], nil
236// }
237//
238// var res []*node
239// for i := 0; i < len(ns); i += 2 {
240// if ns[i] == nil && ns[i+1] == nil {
241// // when both sub nodes are empty, the node is also empty
242// res = append(res, ns[i]) // empty node
243// continue
244// }
245// n := &node{
246// l: ns[i],
247// r: ns[i+1],
248// }
249//
250// if n.l == nil {
251// n.l = &node{
252// h: p.emptyHash,
253// }
254// }
255// if n.r == nil {
256// n.r = &node{
257// h: p.emptyHash,
258// }
259// }
260// if n.l.typ() == vtEmpty && n.r.typ() == vtLeaf {
261// n = n.r
262// }
263// if n.r.typ() == vtEmpty && n.l.typ() == vtLeaf {
264// n = n.l
265// }
266//
267// // once the sub nodes are computed, can compute the current node
268// // hash
269// p.dbg.incHash()
270// k, v, err := newIntermediate(p.hashFunction, n.l.h, n.r.h)
271// if err != nil {
272// return nil, nil, err
273// }
274// n.h = k
275// kv := [2][]byte{k, v}
276// pairs = append(pairs, kv)
277// res = append(res, n)
278// }
279// return upFromNodesComputingHashes(p, res, pairs)
280// }
281
282func (t *vt) add(fromLvl int, k, v []byte) error {
283 leaf := newLeafNode(t.params, k, v)
284 if t.root == nil {
285 t.root = leaf
286 return nil
287 }
288
289 if err := t.root.add(t.params, fromLvl, leaf); err != nil {
290 return err
291 }
292
293 return nil
294}
295
296// computeHashes should be called after all the vt.add is used, once all the
297// leafs are in the tree
298func (t *vt) computeHashes() ([][2][]byte, error) {
299 var err error
300
301 nCPU := flp2(runtime.NumCPU())
302 l := int(math.Log2(float64(nCPU)))
303 nodesAtL, err := t.getNodesAtLevel(l)
304 if err != nil {
305 return nil, err
306 }
307 subRoots := make([]*node, nCPU)
308 bucketPairs := make([][][2][]byte, nCPU)
309 dbgStatsPerBucket := make([]*dbgStats, nCPU)
310
311 var wg sync.WaitGroup
312 wg.Add(nCPU)
313 for i := 0; i < nCPU; i++ {
314 go func(cpu int) {
315 bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
316 bucketVT.params.dbg = newDbgStats()
317 bucketVT.root = nodesAtL[cpu]
318
319 bucketPairs[cpu], err = bucketVT.root.computeHashes(l,
320 t.params.maxLevels, bucketVT.params, bucketPairs[cpu])
321 if err != nil {
322 // TODO WIP
323 fmt.Println("TODO ERR, err:", err)
324 panic(err)
325 }
326
327 subRoots[cpu] = bucketVT.root
328 dbgStatsPerBucket[cpu] = bucketVT.params.dbg
329 wg.Done()
330 }(i)
331 }
332 wg.Wait()
333
334 for i := 0; i < len(dbgStatsPerBucket); i++ {
335 t.params.dbg.add(dbgStatsPerBucket[i])
336 }
337
338 var pairs [][2][]byte
339 for i := 0; i < len(bucketPairs); i++ {
340 pairs = append(pairs, bucketPairs[i]...)
341 }
342
343 nodesAtL, err = t.getNodesAtLevel(l)
344 if err != nil {
345 return nil, err
346 }
347 for i := 0; i < len(nodesAtL); i++ {
348 nodesAtL = subRoots
349 }
350
351 pairs, err = t.root.computeHashes(0, l, t.params, pairs)
352 if err != nil {
353 return nil, err
354 }
355
356 return pairs, nil
357}
358
359func newLeafNode(p *params, k, v []byte) *node {
360 keyPath := make([]byte, p.hashFunction.Len())
361 copy(keyPath[:], k)
362 path := getPath(p.maxLevels, keyPath)
363 n := &node{
364 k: k,
365 v: v,
366 path: path,
367 }
368 return n
369}
370
371type virtualNodeType int
372
373const (
374 vtEmpty = 0 // for convenience uses same value that PrefixValueEmpty
375 vtLeaf = 1 // for convenience uses same value that PrefixValueLeaf
376 vtMid = 2 // for convenience uses same value that PrefixValueIntermediate
377)
378
379func (n *node) typ() virtualNodeType {
380 if n == nil {
381 return vtEmpty // TODO decide if return 'vtEmpty' or an error
382 }
383 if n.l == nil && n.r == nil && n.k != nil {
384 return vtLeaf
385 }
386 if n.l != nil || n.r != nil {
387 return vtMid
388 }
389 return vtEmpty
390}
391
392func (n *node) add(p *params, currLvl int, leaf *node) error {
393 if currLvl > p.maxLevels-1 {
394 return fmt.Errorf("max virtual level %d", currLvl)
395 }
396
397 if n == nil {
398 // n = leaf // TMP!
399 return nil
400 }
401
402 t := n.typ()
403 switch t {
404 case vtMid:
405 if leaf.path[currLvl] {
406 //right
407 if n.r == nil {
408 // empty sub-node, add the leaf here
409 n.r = leaf
410 return nil
411 }
412 if err := n.r.add(p, currLvl+1, leaf); err != nil {
413 return err
414 }
415 } else {
416 if n.l == nil {
417 // empty sub-node, add the leaf here
418 n.l = leaf
419 return nil
420 }
421 if err := n.l.add(p, currLvl+1, leaf); err != nil {
422 return err
423 }
424 }
425 case vtLeaf:
426 if bytes.Equal(n.k, leaf.k) {
427 return fmt.Errorf("key already exists. Existing node: %s, trying to add node: %s",
428 hex.EncodeToString(n.k), hex.EncodeToString(leaf.k))
429 }
430
431 oldLeaf := &node{
432 k: n.k,
433 v: n.v,
434 path: n.path,
435 }
436 // remove values from current node (converting it to mid node)
437 n.k = nil
438 n.v = nil
439 n.h = nil
440 n.path = nil
441 if err := n.downUntilDivergence(p, currLvl, oldLeaf, leaf); err != nil {
442 return err
443 }
444 case vtEmpty:
445 panic(fmt.Errorf("EMPTY %v", n)) // TODO TMP
446 default:
447 return fmt.Errorf("ERR")
448 }
449
450 return nil
451}
452
453func (n *node) downUntilDivergence(p *params, currLvl int, oldLeaf, newLeaf *node) error {
454 if currLvl > p.maxLevels-1 {
455 return fmt.Errorf("max virtual level %d", currLvl)
456 }
457
458 if oldLeaf.path[currLvl] != newLeaf.path[currLvl] {
459 // reached divergence in next level
460 if newLeaf.path[currLvl] {
461 n.l = oldLeaf
462 n.r = newLeaf
463 } else {
464 n.l = newLeaf
465 n.r = oldLeaf
466 }
467 return nil
468 }
469 // no divergence yet, continue going down
470 if newLeaf.path[currLvl] {
471 // right
472 n.r = &node{}
473 if err := n.r.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
474 return err
475 }
476 } else {
477 // left
478 n.l = &node{}
479 if err := n.l.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
480 return err
481 }
482 }
483
484 return nil
485}
486
487// returns an array of key-values to store in the db
488func (n *node) computeHashes(currLvl, maxLvl int, p *params, pairs [][2][]byte) (
489 [][2][]byte, error) {
490 if n == nil || currLvl >= maxLvl {
491 // no need to compute any hash
492 return pairs, nil
493 }
494 if pairs == nil {
495 pairs = [][2][]byte{}
496 }
497 var err error
498 t := n.typ()
499 switch t {
500 case vtLeaf:
501 p.dbg.incHash()
502 leafKey, leafValue, err := newLeafValue(p.hashFunction, n.k, n.v)
503 if err != nil {
504 return pairs, err
505 }
506 n.h = leafKey
507 kv := [2][]byte{leafKey, leafValue}
508 pairs = append(pairs, kv)
509 case vtMid:
510 if n.l != nil {
511 pairs, err = n.l.computeHashes(currLvl+1, maxLvl, p, pairs)
512 if err != nil {
513 return pairs, err
514 }
515 } else {
516 n.l = &node{
517 h: p.emptyHash,
518 }
519 }
520 if n.r != nil {
521 pairs, err = n.r.computeHashes(currLvl+1, maxLvl, p, pairs)
522 if err != nil {
523 return pairs, err
524 }
525 } else {
526 n.r = &node{
527 h: p.emptyHash,
528 }
529 }
530 // once the sub nodes are computed, can compute the current node
531 // hash
532 p.dbg.incHash()
533 k, v, err := newIntermediate(p.hashFunction, n.l.h, n.r.h)
534 if err != nil {
535 return nil, err
536 }
537 n.h = k
538 kv := [2][]byte{k, v}
539 pairs = append(pairs, kv)
540 case vtEmpty:
541 default:
542 fmt.Println("n.computeHashes type no match", t)
543 return nil, fmt.Errorf("ERR TMP") // TODO
544 }
545
546 return pairs, nil
547}
548
549//nolint:unused
550func (t *vt) graphviz(w io.Writer) error {
551 fmt.Fprintf(w, `digraph hierarchy {
552node [fontname=Monospace,fontsize=10,shape=box]
553`)
554 if _, err := t.root.graphviz(w, t.params, 0); err != nil {
555 return err
556 }
557 fmt.Fprintf(w, "}\n")
558 return nil
559}
560
561//nolint:unused
562func (n *node) graphviz(w io.Writer, p *params, nEmpties int) (int, error) {
563 nChars := 4 // TODO move to global constant
564 if n == nil {
565 return nEmpties, nil
566 }
567
568 t := n.typ()
569 switch t {
570 case vtLeaf:
571 leafKey, _, err := newLeafValue(p.hashFunction, n.k, n.v)
572 if err != nil {
573 return nEmpties, err
574 }
575 fmt.Fprintf(w, "\"%p\" [style=filled,label=\"%v\"];\n", n, hex.EncodeToString(leafKey[:nChars]))
576
577 fmt.Fprintf(w, "\"%p\" -> {\"k:%v\\nv:%v\"}\n", n,
578 hex.EncodeToString(n.k[:nChars]),
579 hex.EncodeToString(n.v[:nChars]))
580 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
581 hex.EncodeToString(n.k[:nChars]),
582 hex.EncodeToString(n.v[:nChars]))
583 case vtMid:
584 fmt.Fprintf(w, "\"%p\" [label=\"\"];\n", n)
585
586 lStr := fmt.Sprintf("%p", n.l)
587 rStr := fmt.Sprintf("%p", n.r)
588 eStr := ""
589 if n.l == nil {
590 lStr = fmt.Sprintf("empty%v", nEmpties)
591 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
592 lStr)
593 nEmpties++
594 }
595 if n.r == nil {
596 rStr = fmt.Sprintf("empty%v", nEmpties)
597 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
598 rStr)
599 nEmpties++
600 }
601 fmt.Fprintf(w, "\"%p\" -> {\"%v\" \"%v\"}\n", n, lStr, rStr)
602 fmt.Fprint(w, eStr)
603 nEmpties, err := n.l.graphviz(w, p, nEmpties)
604 if err != nil {
605 return nEmpties, err
606 }
607 nEmpties, err = n.r.graphviz(w, p, nEmpties)
608 if err != nil {
609 return nEmpties, err
610 }
611
612 case vtEmpty:
613 default:
614 return nEmpties, fmt.Errorf("ERR")
615 }
616
617 return nEmpties, nil
618}
619
620//nolint:unused
621func (t *vt) printGraphviz() error {
622 w := bytes.NewBufferString("")
623 fmt.Fprintf(w,
624 "--------\nGraphviz:\n")
625 err := t.graphviz(w)
626 if err != nil {
627 fmt.Println(w)
628 return err
629 }
630 fmt.Fprintf(w,
631 "End of Graphviz --------\n")
632 fmt.Println(w)
633 return nil
634}