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