-
Notifications
You must be signed in to change notification settings - Fork 0
/
counter.go
117 lines (105 loc) · 2.3 KB
/
counter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package radix
import (
"sync/atomic"
"time"
)
type counterPosition struct {
raw uint64
nextChild int
}
// A Counter counts objects in a Tree
type Counter struct {
NodesWithPrefix int
NodeSizes []int
TotalPrefixLength int
TotalPrefixBytes int
LargePrefixes int
Elapsed time.Duration
idx *Tree
stack []counterPosition
}
// NewCounter returns a counter on tree
func NewCounter(tree *Tree) *Counter {
return &Counter{
idx: tree,
NodeSizes: make([]int, 257),
}
}
func (c *Counter) reset() {
c.stack = c.stack[:0]
c.NodesWithPrefix = 0
c.TotalPrefixLength = 0
c.TotalPrefixBytes = 0
c.LargePrefixes = 0
for i := range c.NodeSizes {
c.NodeSizes[i] = 0
}
}
// Count counts nodes and leaves having the given prefix which may be empty
func (c *Counter) Count(searchPrefix []byte) (nodes, leaves int) {
start := time.Now()
c.reset()
alloc := c.idx.GetAllocator()
defer c.idx.ReleaseAllocator(alloc)
alloc.startOp()
defer alloc.endOp()
var p counterPosition
p.raw, _ = c.idx.partialSearch(searchPrefix)
if p.raw == 0 {
return
}
if isLeaf(p.raw) {
leaves++
return
}
nodes++
searchLoop:
for {
_, node, count, prefixLen := explodeNode(p.raw)
data := c.idx.getNodeData(node)
if prefixLen > 0 {
prefixSlots := ((prefixLen + 7) >> 3)
if prefixLen == 255 {
c.LargePrefixes++
prefixLen = int(data[0])
prefixSlots = ((prefixLen + 15) >> 3)
}
c.NodesWithPrefix++
c.TotalPrefixBytes += prefixSlots * 8
c.TotalPrefixLength += prefixLen
data = data[prefixSlots:]
}
c.NodeSizes[count]++
for k := p.nextChild; k < count; k++ {
a := atomic.LoadUint64(&data[k])
if a != 0 {
if isLeaf(a) {
leaves++
} else {
nodes++
next := k + 1
if count >= fullAllocFrom {
// since it's a node-256 we need to check if there are any more non-zero children in this node
for ; next < count; next++ {
if atomic.LoadUint64(&data[next]) != 0 {
break
}
}
}
if next < count {
c.stack = append(c.stack, counterPosition{p.raw, next})
}
p.raw = a
p.nextChild = 0
continue searchLoop
}
}
}
if len(c.stack) == 0 {
break
}
p, c.stack = c.stack[len(c.stack)-1], c.stack[:len(c.stack)-1]
}
c.Elapsed = time.Since(start)
return nodes, leaves
}