fork of https://github.com/sourcegraph/zoekt
1package index
2
3import (
4 "fmt"
5 "testing"
6)
7
8func TestBTree_sorted(t *testing.T) {
9 bt := newBtree(btreeOpts{bucketSize: 2, v: 2})
10 insertMany(t, bt, []ngram{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
11 // inner nodes only
12 //
13 // [3,5,7]
14 // / / \ \
15 // [2] [4] [6] [8,9]
16 //
17 want := "{bucketSize:2 v:2}[3,5,7][2][4][6][8,9]"
18 if s := bt.String(); s != want {
19 t.Fatalf("\nwant:%s\ngot: %s", want, s)
20 }
21}
22
23func TestFindBucket(t *testing.T) {
24 bt := newBtree(btreeOpts{bucketSize: 4, v: 2})
25 insertMany(t, bt, []ngram{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
26
27 buckets := 0
28 offset := 0
29 bt.visit(func(no node) {
30 switch n := no.(type) {
31 case *leaf:
32 n.bucketIndex = buckets
33 buckets++
34 n.postingIndexOffset = offset
35 offset += n.bucketSize
36 case *innerNode:
37 return
38 }
39 })
40
41 cases := []struct {
42 ng ngram
43 wantBucketIndex int
44 wantPostingIndexOffset int
45 }{
46 {
47 ng: 7,
48 wantBucketIndex: 3,
49 wantPostingIndexOffset: 6,
50 },
51 }
52
53 for _, tt := range cases {
54 t.Run(fmt.Sprintf("ngram: %d", tt.ng), func(t *testing.T) {
55 haveBucketIndex, havePostingIndexOffset := bt.find(tt.ng)
56 if tt.wantBucketIndex != haveBucketIndex {
57 t.Fatalf("bucketIndex: want %d, got %d", tt.wantBucketIndex, haveBucketIndex)
58 }
59
60 if tt.wantPostingIndexOffset != havePostingIndexOffset {
61 t.Fatalf("postingIndexOffset: want %d, got %d", tt.wantPostingIndexOffset, havePostingIndexOffset)
62 }
63 })
64 }
65}
66
67func TestGetBucket(t *testing.T) {
68 var off uint32 = 13
69 bucketSize := 4
70
71 cases := []struct {
72 nNgrams int
73 bucketIndex int
74 wantOff uint32
75 wantSz uint32
76 }{
77 // tiny B-tree with just 1 bucket.
78 {
79 nNgrams: 1,
80 bucketIndex: 0,
81 wantOff: off,
82 wantSz: 8,
83 },
84 {
85 nNgrams: 2,
86 bucketIndex: 0,
87 wantOff: off,
88 wantSz: 16,
89 },
90 {
91 nNgrams: 3,
92 bucketIndex: 0,
93 wantOff: off,
94 wantSz: 24,
95 },
96 // B-tree with 10 ngrams, think 1,2,3,4,5,6,7,8,9,10
97 {
98 nNgrams: 10,
99 bucketIndex: 0,
100 wantOff: off,
101 wantSz: 16,
102 },
103 {
104 nNgrams: 10,
105 bucketIndex: 1,
106 wantOff: off + 16,
107 wantSz: 16,
108 },
109 {
110 nNgrams: 10,
111 bucketIndex: 2,
112 wantOff: off + 32,
113 wantSz: 16,
114 },
115 {
116 nNgrams: 10,
117 bucketIndex: 3,
118 wantOff: off + 48,
119 wantSz: 32,
120 },
121 {
122 nNgrams: 9,
123 bucketIndex: 3,
124 wantOff: off + 48,
125 wantSz: 24,
126 },
127 }
128
129 for _, tt := range cases {
130 t.Run("", func(t *testing.T) {
131 bi := btreeIndex{
132 ngramSec: simpleSection{off: off, sz: uint32(tt.nNgrams * ngramEncoding)},
133 }
134
135 bt := newBtree(btreeOpts{
136 bucketSize: bucketSize,
137 v: 2,
138 })
139 for i := range tt.nNgrams {
140 bt.insert(ngram(i + 1))
141 }
142 bt.freeze()
143
144 bi.bt = bt
145
146 off, sz := bi.getBucket(tt.bucketIndex)
147 if off != tt.wantOff {
148 t.Fatalf("off: want %d, got %d", tt.wantOff, off)
149 }
150 if sz != tt.wantSz {
151 t.Fatalf("sz: want %d, got %d", tt.wantSz, sz)
152 }
153 })
154 }
155}
156
157func insertMany(t *testing.T, bt *btree, ngrams []ngram) {
158 t.Helper()
159 for _, ng := range ngrams {
160 bt.insert(ng)
161 }
162}