fork of https://github.com/sourcegraph/zoekt
1// Copyright 2016 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package zoekt
16
17import (
18 "cmp"
19 "encoding/binary"
20 "math"
21 "sort"
22 "unicode"
23 "unicode/utf8"
24)
25
26func generateCaseNgrams(g ngram) []ngram {
27 asRunes := ngramToRunes(g)
28
29 variants := make([]ngram, 0, 8)
30 cur := asRunes
31 for {
32 for i := 0; i < 3; i++ {
33 next := unicode.SimpleFold(cur[i])
34 cur[i] = next
35 if next != asRunes[i] {
36 break
37 }
38 }
39
40 variants = append(variants, runesToNGram(cur))
41 if cur == asRunes {
42 break
43 }
44 }
45
46 return variants
47}
48
49func toLower(in []byte) []byte {
50 out := make([]byte, 0, len(in))
51 var buf [4]byte
52 for _, c := range string(in) {
53 i := utf8.EncodeRune(buf[:], unicode.ToLower(c))
54 out = append(out, buf[:i]...)
55 }
56 return out
57}
58
59// compare 'lower' and 'mixed', where lower is the needle. 'mixed' may
60// be larger than 'lower'. Returns whether there was a match, and if
61// yes, the byte size of the match.
62func caseFoldingEqualsRunes(lower, mixed []byte) (int, bool) {
63 matchTotal := 0
64 for len(lower) > 0 && len(mixed) > 0 {
65 lr, lsz := utf8.DecodeRune(lower)
66 lower = lower[lsz:]
67
68 mr, msz := utf8.DecodeRune(mixed)
69 mixed = mixed[msz:]
70 matchTotal += msz
71
72 if lr != unicode.ToLower(mr) {
73 return 0, false
74 }
75 }
76
77 return matchTotal, len(lower) == 0
78}
79
80type ngram uint64
81
82func runesToNGram(b [ngramSize]rune) ngram {
83 return ngram(uint64(b[0])<<42 | uint64(b[1])<<21 | uint64(b[2]))
84}
85
86func bytesToNGram(b []byte) ngram {
87 return runesToNGram([ngramSize]rune{rune(b[0]), rune(b[1]), rune(b[2])})
88}
89
90func stringToNGram(s string) ngram {
91 return bytesToNGram([]byte(s))
92}
93
94func ngramToBytes(n ngram) []byte {
95 rs := ngramToRunes(n)
96 return []byte{byte(rs[0]), byte(rs[1]), byte(rs[2])}
97}
98
99const runeMask = 1<<21 - 1
100
101func ngramToRunes(n ngram) [ngramSize]rune {
102 return [ngramSize]rune{rune((n >> 42) & runeMask), rune((n >> 21) & runeMask), rune(n & runeMask)}
103}
104
105func (n ngram) String() string {
106 rs := ngramToRunes(n)
107 return string(rs[:])
108}
109
110type runeNgramOff struct {
111 ngram ngram
112 // index is the original index inside of the returned array of splitNGrams
113 index uint32
114}
115
116func (a runeNgramOff) Compare(b runeNgramOff) int {
117 if a.ngram == b.ngram {
118 return cmp.Compare(a.index, b.index)
119 } else if a.ngram < b.ngram {
120 return -1
121 } else {
122 return 1
123 }
124}
125
126func splitNGrams(str []byte) []runeNgramOff {
127 var runeGram [3]rune
128 var off [3]uint32
129 var runeCount int
130
131 result := make([]runeNgramOff, 0, len(str))
132 var i uint32
133
134 for len(str) > 0 {
135 r, sz := utf8.DecodeRune(str)
136 str = str[sz:]
137 runeGram[0] = runeGram[1]
138 off[0] = off[1]
139 runeGram[1] = runeGram[2]
140 off[1] = off[2]
141 runeGram[2] = r
142 off[2] = uint32(i)
143 i += uint32(sz)
144 runeCount++
145 if runeCount < ngramSize {
146 continue
147 }
148
149 ng := runesToNGram(runeGram)
150 result = append(result, runeNgramOff{
151 ngram: ng,
152 index: uint32(len(result)),
153 })
154 }
155 return result
156}
157
158const (
159 _classLowerChar int = iota
160 _classUpperChar
161 _classDigit
162 _classPunct
163 _classOther
164 _classSpace
165)
166
167func byteClass(c byte) int {
168 if c >= 'a' && c <= 'z' {
169 return _classLowerChar
170 }
171 if c >= 'A' && c <= 'Z' {
172 return _classUpperChar
173 }
174 if c >= '0' && c <= '9' {
175 return _classDigit
176 }
177
178 switch c {
179 case ' ', '\n':
180 return _classSpace
181 case '.', ',', ';', '"', '\'':
182 return _classPunct
183 default:
184 return _classOther
185 }
186}
187
188func characterClass(c byte) bool {
189 return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_'
190}
191
192func marshalDocSections(secs []DocumentSection) []byte {
193 ints := make([]uint32, 0, len(secs)*2)
194 for _, s := range secs {
195 ints = append(ints, uint32(s.Start), uint32(s.End))
196 }
197
198 return toSizedDeltas(ints)
199}
200
201func unmarshalDocSections(data []byte, ds []DocumentSection) []DocumentSection {
202 sz, m := binary.Uvarint(data)
203 data = data[m:]
204
205 if cap(ds) < int(sz)/2 {
206 ds = make([]DocumentSection, 0, sz/2)
207 } else {
208 ds = ds[:0]
209 }
210
211 // Inlining the delta decoding to avoid unnecessary allocations that would come
212 // from the straightforward implementation, i.e. packing the result of fromSizedDeltas.
213 var last uint32
214 for len(data) > 0 {
215 var d DocumentSection
216
217 delta, m := binary.Uvarint(data)
218 last += uint32(delta)
219 data = data[m:]
220 d.Start = last
221
222 delta, m = binary.Uvarint(data)
223 last += uint32(delta)
224 data = data[m:]
225 d.End = last
226
227 ds = append(ds, d)
228 }
229 return ds
230}
231
232type ngramSlice []ngram
233
234func (p ngramSlice) Len() int { return len(p) }
235
236func (p ngramSlice) Less(i, j int) bool {
237 return p[i] < p[j]
238}
239
240func (p ngramSlice) Swap(i, j int) {
241 p[i], p[j] = p[j], p[i]
242}
243
244func toSizedDeltas(offsets []uint32) []byte {
245 var enc [8]byte
246
247 deltas := make([]byte, 0, len(offsets)*2)
248
249 m := binary.PutUvarint(enc[:], uint64(len(offsets)))
250 deltas = append(deltas, enc[:m]...)
251
252 var last uint32
253 for _, p := range offsets {
254 delta := p - last
255 last = p
256
257 m := binary.PutUvarint(enc[:], uint64(delta))
258 deltas = append(deltas, enc[:m]...)
259 }
260 return deltas
261}
262
263func fromSizedDeltas(data []byte, ps []uint32) []uint32 {
264 sz, m := binary.Uvarint(data)
265 data = data[m:]
266
267 if cap(ps) < int(sz) {
268 ps = make([]uint32, 0, sz)
269 } else {
270 ps = ps[:0]
271 }
272
273 var last uint32
274 for len(data) > 0 {
275 delta, m := binary.Uvarint(data)
276 offset := last + uint32(delta)
277 last = offset
278 data = data[m:]
279 ps = append(ps, offset)
280 }
281 return ps
282}
283
284func toSizedDeltas16(offsets []uint16) []byte {
285 var enc [8]byte
286
287 deltas := make([]byte, 0, len(offsets)*2)
288
289 m := binary.PutUvarint(enc[:], uint64(len(offsets)))
290 deltas = append(deltas, enc[:m]...)
291
292 var last uint16
293 for _, p := range offsets {
294 delta := p - last
295 last = p
296
297 m := binary.PutUvarint(enc[:], uint64(delta))
298 deltas = append(deltas, enc[:m]...)
299 }
300 return deltas
301}
302
303func fromSizedDeltas16(data []byte, ps []uint16) []uint16 {
304 sz, m := binary.Uvarint(data)
305 data = data[m:]
306
307 if cap(ps) < int(sz) {
308 ps = make([]uint16, 0, sz)
309 } else {
310 ps = ps[:0]
311 }
312
313 var last uint16
314 for len(data) > 0 {
315 delta, m := binary.Uvarint(data)
316 offset := last + uint16(delta)
317 last = offset
318 data = data[m:]
319 ps = append(ps, offset)
320 }
321 return ps
322}
323
324func fromDeltas(data []byte, buf []uint32) []uint32 {
325 buf = buf[:0]
326 if cap(buf) < len(data)/2 {
327 buf = make([]uint32, 0, len(data)/2)
328 }
329
330 var last uint32
331 for len(data) > 0 {
332 delta, m := binary.Uvarint(data)
333 offset := last + uint32(delta)
334 last = offset
335 data = data[m:]
336 buf = append(buf, offset)
337 }
338 return buf
339}
340
341type runeOffsetCorrection struct {
342 runeOffset, byteOffset uint32
343}
344
345// runeOffsetMap converts from rune offsets (with granularity runeOffsetFrequency)
346// to byte offsets, by tracking only the points where a span of runes is non-ASCII,
347// and otherwise interpolating expected byte offsets as one byte per rune.
348//
349// Instead of storing [100, 205, 305], it stores [{x: 200, y: 205}].
350//
351// This is very rarely a slight pessimization on repos where there are frequent
352// non-ASCII characters.
353type runeOffsetMap []runeOffsetCorrection
354
355// makeRuneOffsetMap converts the mostly-predictable runeOffset input
356// into a shorter form tracking the unexpected values.
357//
358// The input is a sequence of y values that we expect to increase by 100 each,
359// so we just store (x, y) points where the expectation is violated.
360func makeRuneOffsetMap(off []uint32) runeOffsetMap {
361 expected := uint32(0)
362 tmp := []runeOffsetCorrection{}
363 for runeOffset, byteOffset := range off {
364 if byteOffset != expected {
365 tmp = append(tmp, runeOffsetCorrection{uint32(runeOffset) * runeOffsetFrequency, byteOffset})
366 expected = byteOffset
367 }
368 expected += runeOffsetFrequency
369 }
370 // copy the slice to ensure it doesn't waste unused trailing capacity
371 out := make([]runeOffsetCorrection, len(tmp))
372 copy(out, tmp)
373 return runeOffsetMap(out)
374}
375
376// lookup converts rune index `off` to a byte offset and a number of additional
377// runes to traverse, given the granularity of runeOffsetFrequency.
378//
379// It does this by finding the nearest point to interpolate from in the map.
380func (m runeOffsetMap) lookup(runeOffset uint32) (uint32, uint32) {
381 left := runeOffset % runeOffsetFrequency
382 runeOffset -= left
383 slen := len(m)
384 if slen == 0 {
385 return runeOffset, left
386 }
387 // sort.Search finds the *first* index for which the predicate is true,
388 // but we want to find the *last* index for which the predicate is true.
389 // This involves some work to reverse the index directions.
390 idx := sort.Search(slen, func(i int) bool {
391 return runeOffset >= m[slen-1-i].runeOffset
392 })
393 idx = slen - 1 - idx
394 // idx is now in the range [-1, len(m))-- -1 indicates that the offset is smaller
395 // than the first entry in the map, so no correction is necessary.
396 byteOff := runeOffset
397 if idx >= 0 {
398 byteOff = m[idx].byteOffset + runeOffset - m[idx].runeOffset
399 }
400 return byteOff, left
401}
402
403func (m runeOffsetMap) sizeBytes() int {
404 return 8 * len(m)
405}
406
407func epsilonEqualsOne(scoreWeight float64) bool {
408 return scoreWeight == 1 || math.Abs(scoreWeight-1.0) < 1e-9
409}