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