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
33type kv struct {
34 pos int // original position in the inputted array
35 keyPath []byte
36 k []byte
37 v []byte
38}
39
40func keysValuesToKvs(maxLevels int, ks, vs [][]byte) ([]kv, []Invalid, error) {
41 if len(ks) != len(vs) {
42 return nil, nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
43 len(ks), len(vs))
44 }
45 var invalids []Invalid
46 var kvs []kv
47 for i := 0; i < len(ks); i++ {
48 keyPath, err := keyPathFromKey(maxLevels, ks[i])
49 if err != nil {
50 invalids = append(invalids, Invalid{i, err})
51 continue
52 }
53 if err := checkKeyValueLen(ks[i], vs[i]); err != nil {
54 invalids = append(invalids, Invalid{i, err})
55 continue
56 }
57
58 var kvsI kv
59 kvsI.pos = i
60 kvsI.keyPath = keyPath
61 kvsI.k = ks[i]
62 kvsI.v = vs[i]
63 kvs = append(kvs, kvsI)
64 }
65
66 return kvs, invalids, nil
67}
68
69// vt stands for virtual tree. It's a tree that does not have any computed hash
70// while placing the leafs. Once all the leafs are placed, it computes all the
71// hashes. In this way, each node hash is only computed one time (at the end)
72// and the tree is computed in memory.
73type vt struct {
74 root *node
75 params *params
76}
77
78func newVT(maxLevels int, hash HashFunction) vt {
79 return vt{
80 root: nil,
81 params: ¶ms{
82 maxLevels: maxLevels,
83 hashFunction: hash,
84 emptyHash: make([]byte, hash.Len()), // empty
85 },
86 }
87}
88
89// addBatch adds a batch of key-values to the VirtualTree. Returns an array
90// containing the indexes of the keys failed to add. Does not include the
91// computation of hashes of the nodes neither the storage of the key-values of
92// the tree into the db. After addBatch, vt.computeHashes should be called to
93// compute the hashes of all the nodes of the tree.
94func (t *vt) addBatch(ks, vs [][]byte) ([]Invalid, error) {
95 nCPU := flp2(runtime.NumCPU())
96 if nCPU == 1 || len(ks) < nCPU {
97 var invalids []Invalid
98 for i := 0; i < len(ks); i++ {
99 if err := t.add(0, ks[i], vs[i]); err != nil {
100 invalids = append(invalids, Invalid{i, err})
101 }
102 }
103 return invalids, nil
104 }
105
106 l := int(math.Log2(float64(nCPU)))
107
108 kvs, invalids, err := keysValuesToKvs(t.params.maxLevels, ks, vs)
109 if err != nil {
110 return invalids, err
111 }
112
113 buckets := splitInBuckets(kvs, nCPU)
114
115 nodesAtL, err := t.getNodesAtLevel(l)
116 if err != nil {
117 return nil, err
118 }
119 if len(nodesAtL) != nCPU && t.root != nil {
120 /*
121 Already populated Tree but Unbalanced
122 - Need to fill M1 and M2, and then will be able to continue with the flow
123 - Search for M1 & M2 in the inputed Keys
124 - Add M1 & M2 to the Tree
125 - From here can continue with the flow
126
127 R
128 / \
129 / \
130 / \
131 * *
132 | \
133 | \
134 | \
135 L: M1 * M2 * (where M1 and M2 are empty)
136 / | /
137 / | /
138 / | /
139 A * *
140 / \ | \
141 / \ | \
142 / \ | \
143 B * * C
144 / \ |\
145 ... ... | \
146 | \
147 D E
148 */
149
150 // add one key at each bucket, and then continue with the flow
151 for i := 0; i < len(buckets); i++ {
152 // add one leaf of the bucket, if there is an error when
153 // adding the k-v, try to add the next one of the bucket
154 // (until one is added)
155 inserted := -1
156 for j := 0; j < len(buckets[i]); j++ {
157 if err := t.add(0, buckets[i][j].k, buckets[i][j].v); err == nil {
158 inserted = j
159 break
160 }
161 }
162
163 // remove the inserted element from buckets[i]
164 if inserted != -1 {
165 buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
166 }
167 }
168 nodesAtL, err = t.getNodesAtLevel(l)
169 if err != nil {
170 return nil, err
171 }
172 }
173 if len(nodesAtL) != nCPU {
174 return nil, fmt.Errorf("This error should not be reached."+
175 " len(nodesAtL) != nCPU, len(nodesAtL)=%d, nCPU=%d."+
176 " Please report it in a new issue:"+
177 " https://github.com/arnaucube/arbo/issues/new", len(nodesAtL), nCPU)
178 }
179
180 subRoots := make([]*node, nCPU)
181 invalidsInBucket := make([][]Invalid, nCPU)
182
183 var wg sync.WaitGroup
184 wg.Add(nCPU)
185 for i := 0; i < nCPU; i++ {
186 go func(cpu int) {
187 bucketVT := newVT(t.params.maxLevels, t.params.hashFunction)
188 bucketVT.root = nodesAtL[cpu]
189 for j := 0; j < len(buckets[cpu]); j++ {
190 if err := bucketVT.add(l, buckets[cpu][j].k,
191 buckets[cpu][j].v); err != nil {
192 invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
193 Invalid{buckets[cpu][j].pos, err})
194 }
195 }
196 subRoots[cpu] = bucketVT.root
197 wg.Done()
198 }(i)
199 }
200 wg.Wait()
201
202 for i := 0; i < len(invalidsInBucket); i++ {
203 invalids = append(invalids, invalidsInBucket[i]...)
204 }
205
206 newRootNode, err := upFromNodes(subRoots)
207 if err != nil {
208 return nil, err
209 }
210 t.root = newRootNode
211
212 return invalids, nil
213}
214
215func (t *vt) getNodesAtLevel(l int) ([]*node, error) {
216 if t.root == nil {
217 var r []*node
218 nChilds := int(math.Pow(2, float64(l))) //nolint:gomnd
219 for i := 0; i < nChilds; i++ {
220 r = append(r, nil)
221 }
222 return r, nil
223 }
224 return t.root.getNodesAtLevel(0, l)
225}
226
227func (n *node) getNodesAtLevel(currLvl, l int) ([]*node, error) {
228 if n == nil {
229 var r []*node
230 nChilds := int(math.Pow(2, float64(l-currLvl))) //nolint:gomnd
231 for i := 0; i < nChilds; i++ {
232 r = append(r, nil)
233 }
234 return r, nil
235 }
236
237 typ := n.typ()
238 if currLvl == l && typ != vtEmpty {
239 return []*node{n}, nil
240 }
241 if currLvl >= l {
242 return nil, fmt.Errorf("This error should not be reached."+
243 " currLvl >= l, currLvl=%d, l=%d."+
244 " Please report it in a new issue:"+
245 " https://github.com/arnaucube/arbo/issues/new", currLvl, l)
246 }
247
248 var nodes []*node
249
250 nodesL, err := n.l.getNodesAtLevel(currLvl+1, l)
251 if err != nil {
252 return nil, err
253 }
254 nodes = append(nodes, nodesL...)
255
256 nodesR, err := n.r.getNodesAtLevel(currLvl+1, l)
257 if err != nil {
258 return nil, err
259 }
260 nodes = append(nodes, nodesR...)
261
262 return nodes, nil
263}
264
265// upFromNodes builds the tree from the bottom to up
266func upFromNodes(ns []*node) (*node, error) {
267 if len(ns) == 1 {
268 return ns[0], nil
269 }
270
271 var res []*node
272 for i := 0; i < len(ns); i += 2 {
273 if (ns[i].typ() == vtEmpty && ns[i+1].typ() == vtEmpty) ||
274 (ns[i].typ() == vtLeaf && ns[i+1].typ() == vtEmpty) {
275 // when both sub nodes are empty, the parent is also empty
276 // or
277 // when 1st sub node is a leaf but the 2nd is empty, the
278 // leaf is used as 'parent'
279 res = append(res, ns[i])
280 continue
281 }
282 if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtLeaf {
283 // when 2nd sub node is a leaf but the 1st is empty, the
284 // leaf is used as 'parent'
285 res = append(res, ns[i+1])
286 continue
287 }
288 n := &node{
289 l: ns[i],
290 r: ns[i+1],
291 }
292 res = append(res, n)
293 }
294 return upFromNodes(res)
295}
296
297// add adds a key&value as a leaf in the VirtualTree
298func (t *vt) add(fromLvl int, k, v []byte) error {
299 leaf, err := newLeafNode(t.params, k, v)
300 if err != nil {
301 return err
302 }
303 if t.root == nil {
304 t.root = leaf
305 return nil
306 }
307
308 if err := t.root.add(t.params, fromLvl, leaf); err != nil {
309 return err
310 }
311
312 return nil
313}
314
315// computeHashes should be called after all the vt.add is used, once all the
316// leafs are in the tree. Computes the hashes of the tree, parallelizing in the
317// available CPUs.
318func (t *vt) computeHashes() ([][2][]byte, error) {
319 var err error
320
321 nCPU := flp2(runtime.NumCPU())
322 l := int(math.Log2(float64(nCPU)))
323 nodesAtL, err := t.getNodesAtLevel(l)
324 if err != nil {
325 return nil, err
326 }
327 subRoots := make([]*node, nCPU)
328 bucketPairs := make([][][2][]byte, nCPU)
329 dbgStatsPerBucket := make([]*dbgStats, nCPU)
330 errs := make([]error, nCPU)
331
332 var wg sync.WaitGroup
333 wg.Add(nCPU)
334 for i := 0; i < nCPU; i++ {
335 go func(cpu int) {
336 bucketVT := newVT(t.params.maxLevels, t.params.hashFunction)
337 bucketVT.params.dbg = newDbgStats()
338 bucketVT.root = nodesAtL[cpu]
339 var err error
340 bucketPairs[cpu], err = bucketVT.root.computeHashes(l-1,
341 t.params.maxLevels, bucketVT.params, bucketPairs[cpu])
342 if err != nil {
343 errs[cpu] = err
344 }
345
346 subRoots[cpu] = bucketVT.root
347 dbgStatsPerBucket[cpu] = bucketVT.params.dbg
348 wg.Done()
349 }(i)
350 }
351 wg.Wait()
352
353 for i := 0; i < len(errs); i++ {
354 if errs[i] != nil {
355 return nil, errs[i]
356 }
357 }
358
359 for i := 0; i < len(dbgStatsPerBucket); i++ {
360 t.params.dbg.add(dbgStatsPerBucket[i])
361 }
362
363 var pairs [][2][]byte
364 for i := 0; i < len(bucketPairs); i++ {
365 pairs = append(pairs, bucketPairs[i]...)
366 }
367
368 nodesAtL, err = t.getNodesAtLevel(l)
369 if err != nil {
370 return nil, err
371 }
372 for i := 0; i < len(nodesAtL); i++ {
373 nodesAtL = subRoots
374 }
375
376 pairs, err = t.root.computeHashes(0, l, t.params, pairs)
377 if err != nil {
378 return nil, err
379 }
380
381 return pairs, nil
382}
383
384func newLeafNode(p *params, k, v []byte) (*node, error) {
385 if err := checkKeyValueLen(k, v); err != nil {
386 return nil, err
387 }
388
389 keyPath, err := keyPathFromKey(p.maxLevels, k)
390 if err != nil {
391 return nil, err
392 }
393 path := getPath(p.maxLevels, keyPath)
394 n := &node{
395 k: k,
396 v: v,
397 path: path,
398 }
399 return n, nil
400}
401
402type virtualNodeType int
403
404const (
405 vtEmpty = 0 // for convenience uses same value that PrefixValueEmpty
406 vtLeaf = 1 // for convenience uses same value that PrefixValueLeaf
407 vtMid = 2 // for convenience uses same value that PrefixValueIntermediate
408)
409
410func (n *node) typ() virtualNodeType {
411 if n == nil {
412 return vtEmpty
413 }
414 if n.l == nil && n.r == nil && n.k != nil {
415 return vtLeaf
416 }
417 if n.l != nil || n.r != nil {
418 return vtMid
419 }
420 return vtEmpty
421}
422
423func (n *node) add(p *params, currLvl int, leaf *node) error {
424 if currLvl > p.maxLevels-1 {
425 return ErrMaxVirtualLevel
426 }
427
428 if n == nil {
429 // n = leaf // TMP!
430 return nil
431 }
432
433 t := n.typ()
434 switch t {
435 case vtMid:
436 if leaf.path[currLvl] {
437 //right
438 if n.r == nil {
439 // empty sub-node, add the leaf here
440 n.r = leaf
441 return nil
442 }
443 if err := n.r.add(p, currLvl+1, leaf); err != nil {
444 return err
445 }
446 } else {
447 if n.l == nil {
448 // empty sub-node, add the leaf here
449 n.l = leaf
450 return nil
451 }
452 if err := n.l.add(p, currLvl+1, leaf); err != nil {
453 return err
454 }
455 }
456 case vtLeaf:
457 if bytes.Equal(n.k, leaf.k) {
458 return fmt.Errorf("%s. Existing node: %s, trying to add node: %s",
459 ErrKeyAlreadyExists, hex.EncodeToString(n.k),
460 hex.EncodeToString(leaf.k))
461 }
462
463 oldLeaf := &node{
464 k: n.k,
465 v: n.v,
466 path: n.path,
467 }
468 // remove values from current node (converting it to mid node)
469 n.k = nil
470 n.v = nil
471 n.h = nil
472 n.path = nil
473 if err := n.downUntilDivergence(p, currLvl, oldLeaf, leaf); err != nil {
474 return err
475 }
476 case vtEmpty:
477 return fmt.Errorf("virtual tree node.add() with empty node %v", n)
478 default:
479 return fmt.Errorf("virtual tree node.add() with unknown node type %v", n)
480 }
481
482 return nil
483}
484
485func (n *node) downUntilDivergence(p *params, currLvl int, oldLeaf, newLeaf *node) error {
486 if currLvl > p.maxLevels-1 {
487 return ErrMaxVirtualLevel
488 }
489
490 if oldLeaf.path[currLvl] != newLeaf.path[currLvl] {
491 // reached divergence in next level
492 if newLeaf.path[currLvl] {
493 n.l = oldLeaf
494 n.r = newLeaf
495 } else {
496 n.l = newLeaf
497 n.r = oldLeaf
498 }
499 return nil
500 }
501 // no divergence yet, continue going down
502 if newLeaf.path[currLvl] {
503 // right
504 n.r = &node{}
505 if err := n.r.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
506 return err
507 }
508 } else {
509 // left
510 n.l = &node{}
511 if err := n.l.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
512 return err
513 }
514 }
515
516 return nil
517}
518
519func splitInBuckets(kvs []kv, nBuckets int) [][]kv {
520 buckets := make([][]kv, nBuckets)
521 // 1. classify the keyvalues into buckets
522 for i := 0; i < len(kvs); i++ {
523 pair := kvs[i]
524
525 // bucketnum := keyToBucket(pair.k, nBuckets)
526 bucketnum := keyToBucket(pair.keyPath, nBuckets)
527 buckets[bucketnum] = append(buckets[bucketnum], pair)
528 }
529 return buckets
530}
531
532// TODO rename in a more 'real' name (calculate bucket from/for key)
533func keyToBucket(k []byte, nBuckets int) int {
534 nLevels := int(math.Log2(float64(nBuckets)))
535 b := make([]int, nBuckets)
536 for i := 0; i < nBuckets; i++ {
537 b[i] = i
538 }
539 r := b
540 mid := len(r) / 2 //nolint:gomnd
541 for i := 0; i < nLevels; i++ {
542 if int(k[i/bitsPerByte]&(1<<(i%bitsPerByte))) != 0 {
543 r = r[mid:]
544 mid = len(r) / 2 //nolint:gomnd
545 } else {
546 r = r[:mid]
547 mid = len(r) / 2 //nolint:gomnd
548 }
549 }
550 return r[0]
551}
552
553// flp2 computes the floor power of 2, the highest power of 2 under the given
554// value.
555func flp2(n int) int {
556 res := 0
557 for i := n; i >= 1; i-- {
558 if (i & (i - 1)) == 0 {
559 res = i
560 break
561 }
562 }
563 return res
564}
565
566// computeHashes computes the hashes under the node from which is called the
567// method. Returns an array of key-values to store in the db
568func (n *node) computeHashes(currLvl, maxLvl int, p *params, pairs [][2][]byte) (
569 [][2][]byte, error) {
570 if n == nil || currLvl >= maxLvl {
571 // no need to compute any hash
572 return pairs, nil
573 }
574 if pairs == nil {
575 pairs = [][2][]byte{}
576 }
577 var err error
578 t := n.typ()
579 switch t {
580 case vtLeaf:
581 p.dbg.incHash()
582 leafKey, leafValue, err := newLeafValue(p.hashFunction, n.k, n.v)
583 if err != nil {
584 return pairs, err
585 }
586 n.h = leafKey
587 kv := [2][]byte{leafKey, leafValue}
588 pairs = append(pairs, kv)
589 case vtMid:
590 if n.l != nil {
591 pairs, err = n.l.computeHashes(currLvl+1, maxLvl, p, pairs)
592 if err != nil {
593 return pairs, err
594 }
595 } else {
596 n.l = &node{
597 h: p.emptyHash,
598 }
599 }
600 if n.r != nil {
601 pairs, err = n.r.computeHashes(currLvl+1, maxLvl, p, pairs)
602 if err != nil {
603 return pairs, err
604 }
605 } else {
606 n.r = &node{
607 h: p.emptyHash,
608 }
609 }
610 // once the sub nodes are computed, can compute the current node
611 // hash
612 p.dbg.incHash()
613 k, v, err := newIntermediate(p.hashFunction, n.l.h, n.r.h)
614 if err != nil {
615 return nil, err
616 }
617 n.h = k
618 kv := [2][]byte{k, v}
619 pairs = append(pairs, kv)
620 case vtEmpty:
621 default:
622 return nil, fmt.Errorf("error: n.computeHashes type (%d) no match", t)
623 }
624
625 return pairs, nil
626}
627
628//nolint:unused
629func (t *vt) graphviz(w io.Writer) error {
630 fmt.Fprintf(w, `digraph hierarchy {
631node [fontname=Monospace,fontsize=10,shape=box]
632`)
633 if _, err := t.root.graphviz(w, t.params, 0); err != nil {
634 return err
635 }
636 fmt.Fprintf(w, "}\n")
637 return nil
638}
639
640//nolint:unused
641func (n *node) graphviz(w io.Writer, p *params, nEmpties int) (int, error) {
642 if n == nil {
643 return nEmpties, nil
644 }
645
646 t := n.typ()
647 switch t {
648 case vtLeaf:
649 leafKey, _, err := newLeafValue(p.hashFunction, n.k, n.v)
650 if err != nil {
651 return nEmpties, err
652 }
653 fmt.Fprintf(w, "\"%p\" [style=filled,label=\"%v\"];\n", n, hex.EncodeToString(leafKey[:nChars]))
654
655 k := n.k
656 v := n.v
657 if len(n.k) >= nChars {
658 k = n.k[:nChars]
659 }
660 if len(n.v) >= nChars {
661 v = n.v[:nChars]
662 }
663
664 fmt.Fprintf(w, "\"%p\" -> {\"k:%v\\nv:%v\"}\n", n,
665 hex.EncodeToString(k),
666 hex.EncodeToString(v))
667 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
668 hex.EncodeToString(k),
669 hex.EncodeToString(v))
670 case vtMid:
671 fmt.Fprintf(w, "\"%p\" [label=\"\"];\n", n)
672
673 lStr := fmt.Sprintf("%p", n.l)
674 rStr := fmt.Sprintf("%p", n.r)
675 eStr := ""
676 if n.l == nil {
677 lStr = fmt.Sprintf("empty%v", nEmpties)
678 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
679 lStr)
680 nEmpties++
681 }
682 if n.r == nil {
683 rStr = fmt.Sprintf("empty%v", nEmpties)
684 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
685 rStr)
686 nEmpties++
687 }
688 fmt.Fprintf(w, "\"%p\" -> {\"%v\" \"%v\"}\n", n, lStr, rStr)
689 fmt.Fprint(w, eStr)
690 nEmpties, err := n.l.graphviz(w, p, nEmpties)
691 if err != nil {
692 return nEmpties, err
693 }
694 nEmpties, err = n.r.graphviz(w, p, nEmpties)
695 if err != nil {
696 return nEmpties, err
697 }
698
699 case vtEmpty:
700 default:
701 return nEmpties, fmt.Errorf("ERR")
702 }
703
704 return nEmpties, nil
705}
706
707//nolint:unused
708func (t *vt) printGraphviz() error {
709 w := bytes.NewBufferString("")
710 fmt.Fprintf(w,
711 "--------\nGraphviz:\n")
712 err := t.graphviz(w)
713 if err != nil {
714 fmt.Println(w)
715 return err
716 }
717 fmt.Fprintf(w,
718 "End of Graphviz --------\n")
719 fmt.Println(w)
720 return nil
721}