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