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 panic("should not happen") // TODO TMP
164 }
165
166 subRoots := make([]*node, nCPU)
167 invalidsInBucket := make([][]int, nCPU)
168
169 var wg sync.WaitGroup
170 wg.Add(nCPU)
171 for i := 0; i < nCPU; i++ {
172 go func(cpu int) {
173 bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
174 bucketVT.root = nodesAtL[cpu]
175 for j := 0; j < len(buckets[cpu]); j++ {
176 if err = bucketVT.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
177 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
178 }
179 }
180 subRoots[cpu] = bucketVT.root
181 wg.Done()
182 }(i)
183 }
184 wg.Wait()
185
186 var invalids []int
187 for i := 0; i < len(invalidsInBucket); i++ {
188 invalids = append(invalids, invalidsInBucket[i]...)
189 }
190
191 newRootNode, err := upFromNodes(subRoots)
192 if err != nil {
193 return nil, err
194 }
195 t.root = newRootNode
196
197 return invalids, nil
198}
199
200func (t *vt) getNodesAtLevel(l int) ([]*node, error) {
201 if t.root == nil {
202 var r []*node
203 nChilds := int(math.Pow(2, float64(l))) //nolint:gomnd
204 for i := 0; i < nChilds; i++ {
205 r = append(r, nil)
206 }
207 return r, nil
208 }
209 return t.root.getNodesAtLevel(0, l)
210}
211
212func (n *node) getNodesAtLevel(currLvl, l int) ([]*node, error) {
213 if n == nil {
214 var r []*node
215 nChilds := int(math.Pow(2, float64(l-currLvl))) //nolint:gomnd
216 for i := 0; i < nChilds; i++ {
217 r = append(r, nil)
218 }
219 return r, nil
220 }
221
222 typ := n.typ()
223 if currLvl == l && typ != vtEmpty {
224 return []*node{n}, nil
225 }
226 if currLvl >= l {
227 panic("should not reach this point") // TODO TMP
228 }
229
230 var nodes []*node
231
232 nodesL, err := n.l.getNodesAtLevel(currLvl+1, l)
233 if err != nil {
234 return nil, err
235 }
236 nodes = append(nodes, nodesL...)
237
238 nodesR, err := n.r.getNodesAtLevel(currLvl+1, l)
239 if err != nil {
240 return nil, err
241 }
242 nodes = append(nodes, nodesR...)
243
244 return nodes, nil
245}
246
247// upFromNodes builds the tree from the bottom to up
248func upFromNodes(ns []*node) (*node, error) {
249 if len(ns) == 1 {
250 return ns[0], nil
251 }
252
253 var res []*node
254 for i := 0; i < len(ns); i += 2 {
255 if (ns[i].typ() == vtEmpty && ns[i+1].typ() == vtEmpty) ||
256 (ns[i].typ() == vtLeaf && ns[i+1].typ() == vtEmpty) {
257 // when both sub nodes are empty, the parent is also empty
258 // or
259 // when 1st sub node is a leaf but the 2nd is empty, the
260 // leaf is used as parent
261 res = append(res, ns[i])
262 continue
263 }
264 if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtLeaf {
265 // when 2nd sub node is a leaf but the 1st is empty, the
266 // leaf is used as 'parent'
267 res = append(res, ns[i+1])
268 continue
269 }
270 n := &node{
271 l: ns[i],
272 r: ns[i+1],
273 }
274 res = append(res, n)
275 }
276 return upFromNodes(res)
277}
278
279// add adds a key&value as a leaf in the VirtualTree
280func (t *vt) add(fromLvl int, k, v []byte) error {
281 leaf := newLeafNode(t.params, k, v)
282 if t.root == nil {
283 t.root = leaf
284 return nil
285 }
286
287 if err := t.root.add(t.params, fromLvl, leaf); err != nil {
288 return err
289 }
290
291 return nil
292}
293
294// computeHashes should be called after all the vt.add is used, once all the
295// leafs are in the tree. Computes the hashes of the tree, parallelizing in the
296// available CPUs.
297func (t *vt) computeHashes() ([][2][]byte, error) {
298 var err error
299
300 nCPU := flp2(runtime.NumCPU())
301 l := int(math.Log2(float64(nCPU)))
302 nodesAtL, err := t.getNodesAtLevel(l)
303 if err != nil {
304 return nil, err
305 }
306 subRoots := make([]*node, nCPU)
307 bucketPairs := make([][][2][]byte, nCPU)
308 dbgStatsPerBucket := make([]*dbgStats, nCPU)
309
310 var wg sync.WaitGroup
311 wg.Add(nCPU)
312 for i := 0; i < nCPU; i++ {
313 go func(cpu int) {
314 bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
315 bucketVT.params.dbg = newDbgStats()
316 bucketVT.root = nodesAtL[cpu]
317
318 bucketPairs[cpu], err = bucketVT.root.computeHashes(l,
319 t.params.maxLevels, bucketVT.params, bucketPairs[cpu])
320 if err != nil {
321 // TODO WIP
322 panic("TODO: " + err.Error())
323 }
324
325 subRoots[cpu] = bucketVT.root
326 dbgStatsPerBucket[cpu] = bucketVT.params.dbg
327 wg.Done()
328 }(i)
329 }
330 wg.Wait()
331
332 for i := 0; i < len(dbgStatsPerBucket); i++ {
333 t.params.dbg.add(dbgStatsPerBucket[i])
334 }
335
336 var pairs [][2][]byte
337 for i := 0; i < len(bucketPairs); i++ {
338 pairs = append(pairs, bucketPairs[i]...)
339 }
340
341 nodesAtL, err = t.getNodesAtLevel(l)
342 if err != nil {
343 return nil, err
344 }
345 for i := 0; i < len(nodesAtL); i++ {
346 nodesAtL = subRoots
347 }
348
349 pairs, err = t.root.computeHashes(0, l, t.params, pairs)
350 if err != nil {
351 return nil, err
352 }
353
354 return pairs, nil
355}
356
357func newLeafNode(p *params, k, v []byte) *node {
358 keyPath := make([]byte, p.hashFunction.Len())
359 copy(keyPath[:], k)
360 path := getPath(p.maxLevels, keyPath)
361 n := &node{
362 k: k,
363 v: v,
364 path: path,
365 }
366 return n
367}
368
369type virtualNodeType int
370
371const (
372 vtEmpty = 0 // for convenience uses same value that PrefixValueEmpty
373 vtLeaf = 1 // for convenience uses same value that PrefixValueLeaf
374 vtMid = 2 // for convenience uses same value that PrefixValueIntermediate
375)
376
377func (n *node) typ() virtualNodeType {
378 if n == nil {
379 return vtEmpty // TODO decide if return 'vtEmpty' or an error
380 }
381 if n.l == nil && n.r == nil && n.k != nil {
382 return vtLeaf
383 }
384 if n.l != nil || n.r != nil {
385 return vtMid
386 }
387 return vtEmpty
388}
389
390func (n *node) add(p *params, currLvl int, leaf *node) error {
391 if currLvl > p.maxLevels-1 {
392 return ErrMaxVirtualLevel
393 }
394
395 if n == nil {
396 // n = leaf // TMP!
397 return nil
398 }
399
400 t := n.typ()
401 switch t {
402 case vtMid:
403 if leaf.path[currLvl] {
404 //right
405 if n.r == nil {
406 // empty sub-node, add the leaf here
407 n.r = leaf
408 return nil
409 }
410 if err := n.r.add(p, currLvl+1, leaf); err != nil {
411 return err
412 }
413 } else {
414 if n.l == nil {
415 // empty sub-node, add the leaf here
416 n.l = leaf
417 return nil
418 }
419 if err := n.l.add(p, currLvl+1, leaf); err != nil {
420 return err
421 }
422 }
423 case vtLeaf:
424 if bytes.Equal(n.k, leaf.k) {
425 return fmt.Errorf("%s. Existing node: %s, trying to add node: %s",
426 ErrKeyAlreadyExists, hex.EncodeToString(n.k),
427 hex.EncodeToString(leaf.k))
428 }
429
430 oldLeaf := &node{
431 k: n.k,
432 v: n.v,
433 path: n.path,
434 }
435 // remove values from current node (converting it to mid node)
436 n.k = nil
437 n.v = nil
438 n.h = nil
439 n.path = nil
440 if err := n.downUntilDivergence(p, currLvl, oldLeaf, leaf); err != nil {
441 return err
442 }
443 case vtEmpty:
444 panic(fmt.Errorf("EMPTY %v", n)) // TODO TMP
445 default:
446 return fmt.Errorf("ERR") // TODO TMP
447 }
448
449 return nil
450}
451
452func (n *node) downUntilDivergence(p *params, currLvl int, oldLeaf, newLeaf *node) error {
453 if currLvl > p.maxLevels-1 {
454 return ErrMaxVirtualLevel
455 }
456
457 if oldLeaf.path[currLvl] != newLeaf.path[currLvl] {
458 // reached divergence in next level
459 if newLeaf.path[currLvl] {
460 n.l = oldLeaf
461 n.r = newLeaf
462 } else {
463 n.l = newLeaf
464 n.r = oldLeaf
465 }
466 return nil
467 }
468 // no divergence yet, continue going down
469 if newLeaf.path[currLvl] {
470 // right
471 n.r = &node{}
472 if err := n.r.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
473 return err
474 }
475 } else {
476 // left
477 n.l = &node{}
478 if err := n.l.downUntilDivergence(p, currLvl+1, oldLeaf, newLeaf); err != nil {
479 return err
480 }
481 }
482
483 return nil
484}
485
486func splitInBuckets(kvs []kv, nBuckets int) [][]kv {
487 buckets := make([][]kv, nBuckets)
488 // 1. classify the keyvalues into buckets
489 for i := 0; i < len(kvs); i++ {
490 pair := kvs[i]
491
492 // bucketnum := keyToBucket(pair.k, nBuckets)
493 bucketnum := keyToBucket(pair.keyPath, nBuckets)
494 buckets[bucketnum] = append(buckets[bucketnum], pair)
495 }
496 return buckets
497}
498
499// TODO rename in a more 'real' name (calculate bucket from/for key)
500func keyToBucket(k []byte, nBuckets int) int {
501 nLevels := int(math.Log2(float64(nBuckets)))
502 b := make([]int, nBuckets)
503 for i := 0; i < nBuckets; i++ {
504 b[i] = i
505 }
506 r := b
507 mid := len(r) / 2 //nolint:gomnd
508 for i := 0; i < nLevels; i++ {
509 if int(k[i/8]&(1<<(i%8))) != 0 {
510 r = r[mid:]
511 mid = len(r) / 2 //nolint:gomnd
512 } else {
513 r = r[:mid]
514 mid = len(r) / 2 //nolint:gomnd
515 }
516 }
517 return r[0]
518}
519
520// flp2 computes the floor power of 2, the highest power of 2 under the given
521// value.
522func flp2(n int) int {
523 res := 0
524 for i := n; i >= 1; i-- {
525 if (i & (i - 1)) == 0 {
526 res = i
527 break
528 }
529 }
530 return res
531}
532
533// computeHashes computes the hashes under the node from which is called the
534// method. Returns an array of key-values to store in the db
535func (n *node) computeHashes(currLvl, maxLvl int, p *params, pairs [][2][]byte) (
536 [][2][]byte, error) {
537 if n == nil || currLvl >= maxLvl {
538 // no need to compute any hash
539 return pairs, nil
540 }
541 if pairs == nil {
542 pairs = [][2][]byte{}
543 }
544 var err error
545 t := n.typ()
546 switch t {
547 case vtLeaf:
548 p.dbg.incHash()
549 leafKey, leafValue, err := newLeafValue(p.hashFunction, n.k, n.v)
550 if err != nil {
551 return pairs, err
552 }
553 n.h = leafKey
554 kv := [2][]byte{leafKey, leafValue}
555 pairs = append(pairs, kv)
556 case vtMid:
557 if n.l != nil {
558 pairs, err = n.l.computeHashes(currLvl+1, maxLvl, p, pairs)
559 if err != nil {
560 return pairs, err
561 }
562 } else {
563 n.l = &node{
564 h: p.emptyHash,
565 }
566 }
567 if n.r != nil {
568 pairs, err = n.r.computeHashes(currLvl+1, maxLvl, p, pairs)
569 if err != nil {
570 return pairs, err
571 }
572 } else {
573 n.r = &node{
574 h: p.emptyHash,
575 }
576 }
577 // once the sub nodes are computed, can compute the current node
578 // hash
579 p.dbg.incHash()
580 k, v, err := newIntermediate(p.hashFunction, n.l.h, n.r.h)
581 if err != nil {
582 return nil, err
583 }
584 n.h = k
585 kv := [2][]byte{k, v}
586 pairs = append(pairs, kv)
587 case vtEmpty:
588 default:
589 return nil, fmt.Errorf("ERR:n.computeHashes type (%d) no match", t) // TODO TMP
590 }
591
592 return pairs, nil
593}
594
595//nolint:unused
596func (t *vt) graphviz(w io.Writer) error {
597 fmt.Fprintf(w, `digraph hierarchy {
598node [fontname=Monospace,fontsize=10,shape=box]
599`)
600 if _, err := t.root.graphviz(w, t.params, 0); err != nil {
601 return err
602 }
603 fmt.Fprintf(w, "}\n")
604 return nil
605}
606
607//nolint:unused
608func (n *node) graphviz(w io.Writer, p *params, nEmpties int) (int, error) {
609 if n == nil {
610 return nEmpties, nil
611 }
612
613 t := n.typ()
614 switch t {
615 case vtLeaf:
616 leafKey, _, err := newLeafValue(p.hashFunction, n.k, n.v)
617 if err != nil {
618 return nEmpties, err
619 }
620 fmt.Fprintf(w, "\"%p\" [style=filled,label=\"%v\"];\n", n, hex.EncodeToString(leafKey[:nChars]))
621
622 k := n.k
623 v := n.v
624 if len(n.k) >= nChars {
625 k = n.k[:nChars]
626 }
627 if len(n.v) >= nChars {
628 v = n.v[:nChars]
629 }
630
631 fmt.Fprintf(w, "\"%p\" -> {\"k:%v\\nv:%v\"}\n", n,
632 hex.EncodeToString(k),
633 hex.EncodeToString(v))
634 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
635 hex.EncodeToString(k),
636 hex.EncodeToString(v))
637 case vtMid:
638 fmt.Fprintf(w, "\"%p\" [label=\"\"];\n", n)
639
640 lStr := fmt.Sprintf("%p", n.l)
641 rStr := fmt.Sprintf("%p", n.r)
642 eStr := ""
643 if n.l == nil {
644 lStr = fmt.Sprintf("empty%v", nEmpties)
645 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
646 lStr)
647 nEmpties++
648 }
649 if n.r == nil {
650 rStr = fmt.Sprintf("empty%v", nEmpties)
651 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
652 rStr)
653 nEmpties++
654 }
655 fmt.Fprintf(w, "\"%p\" -> {\"%v\" \"%v\"}\n", n, lStr, rStr)
656 fmt.Fprint(w, eStr)
657 nEmpties, err := n.l.graphviz(w, p, nEmpties)
658 if err != nil {
659 return nEmpties, err
660 }
661 nEmpties, err = n.r.graphviz(w, p, nEmpties)
662 if err != nil {
663 return nEmpties, err
664 }
665
666 case vtEmpty:
667 default:
668 return nEmpties, fmt.Errorf("ERR")
669 }
670
671 return nEmpties, nil
672}
673
674//nolint:unused
675func (t *vt) printGraphviz() error {
676 w := bytes.NewBufferString("")
677 fmt.Fprintf(w,
678 "--------\nGraphviz:\n")
679 err := t.graphviz(w)
680 if err != nil {
681 fmt.Println(w)
682 return err
683 }
684 fmt.Fprintf(w,
685 "End of Graphviz --------\n")
686 fmt.Println(w)
687 return nil
688}