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