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 "encoding/binary"
19 "errors"
20 "fmt"
21 "hash/crc64"
22 "log"
23 "math/bits"
24 "slices"
25 "unicode/utf8"
26
27 "github.com/sourcegraph/zoekt/query"
28)
29
30// indexData holds the pattern-independent data that we have to have
31// in memory to search. Most of the memory is taken up by the ngram =>
32// offset index.
33type indexData struct {
34 symbols symbolData
35
36 file IndexFile
37
38 contentNgrams btreeIndex
39
40 newlinesStart uint32
41 newlinesIndex []uint32
42
43 docSectionsStart uint32
44 docSectionsIndex []uint32
45
46 runeDocSections []DocumentSection
47
48 // rune offset=>byte offset mapping, relative to the start of the content corpus
49 runeOffsets runeOffsetMap
50
51 // offsets of file contents; includes end of last file
52 boundariesStart uint32
53 boundaries []uint32
54
55 // rune offsets for the file content boundaries
56 fileEndRunes []uint32
57
58 fileNameContent []byte
59 fileNameIndex []uint32
60 fileNameNgrams btreeIndex
61
62 // fileEndSymbol[i] is the index of the first symbol for document i.
63 fileEndSymbol []uint32
64
65 // rune offset=>byte offset mapping, relative to the start of the filename corpus
66 fileNameRuneOffsets runeOffsetMap
67
68 // rune offsets for the file name boundaries
69 fileNameEndRunes []uint32
70
71 fileBranchMasks []uint64
72
73 // mask (power of 2) => name
74 branchNames []map[uint]string
75
76 // name => mask (power of 2)
77 branchIDs []map[string]uint
78
79 metaData IndexMetadata
80 repoMetaData []Repository
81
82 subRepos []uint32
83 subRepoPaths [][]string
84
85 // Checksums for all the files, at 8-byte intervals
86 checksums []byte
87
88 // languages for all the files.
89 languages []byte
90
91 // inverse of LanguageMap in metaData
92 languageMap map[uint16]string
93
94 repoListEntry []RepoListEntry
95
96 // repository indexes for all the files
97 repos []uint16
98
99 // Experimental: docID => rank vec
100 ranks [][]float64
101
102 // rawConfigMasks contains the encoded RawConfig for each repository
103 rawConfigMasks []uint8
104}
105
106type symbolData struct {
107 // symContent stores Symbol.Sym and Symbol.Parent.
108 // TODO we don't need to store Symbol.Sym.
109 symContent []byte
110 symIndex []byte
111 // symKindContent is an enum of sym.Kind and sym.ParentKind
112 symKindContent []byte
113 symKindIndex []uint32
114 // symMetadata is [4]uint32 0 Kind Parent ParentKind
115 symMetaData []byte
116}
117
118func uint32SliceAt(a []byte, n uint32) uint32 {
119 return binary.BigEndian.Uint32(a[n*4:])
120}
121
122func uint32SliceLen(a []byte) uint32 {
123 return uint32(len(a) / 4)
124}
125
126// parent returns index i of the parent enum
127func (d *symbolData) parent(i uint32) []byte {
128 delta := uint32SliceAt(d.symIndex, 0)
129 start := uint32SliceAt(d.symIndex, i) - delta
130 var end uint32
131 if i+1 == uint32SliceLen(d.symIndex) {
132 end = uint32(len(d.symContent))
133 } else {
134 end = uint32SliceAt(d.symIndex, i+1) - delta
135 }
136 return d.symContent[start:end]
137}
138
139// kind returns index i of the kind enum
140func (d *symbolData) kind(i uint32) []byte {
141 return d.symKindContent[d.symKindIndex[i]:d.symKindIndex[i+1]]
142}
143
144// data returns the symbol at index i
145func (d *symbolData) data(i uint32) *Symbol {
146 size := uint32(4 * 4) // 4 uint32s
147 offset := i * size
148 if offset >= uint32(len(d.symMetaData)) {
149 return nil
150 }
151
152 metadata := d.symMetaData[offset : offset+size]
153 sym := &Symbol{}
154 key := uint32SliceAt(metadata, 1)
155 sym.Kind = string(d.kind(key))
156 key = uint32SliceAt(metadata, 2)
157 sym.Parent = string(d.parent(key))
158 key = uint32SliceAt(metadata, 3)
159 sym.ParentKind = string(d.kind(key))
160 return sym
161}
162
163func (d *indexData) getChecksum(idx uint32) []byte {
164 start := crc64.Size * idx
165 return d.checksums[start : start+crc64.Size]
166}
167
168func (d *indexData) getLanguage(idx uint32) uint16 {
169 if d.metaData.IndexFeatureVersion < 12 {
170 // older zoekt files had 8-bit language entries
171 return uint16(d.languages[idx])
172 }
173 // newer zoekt files have 16-bit language entries
174 return uint16(d.languages[idx*2]) | uint16(d.languages[idx*2+1])<<8
175}
176
177// calculates stats for files in the range [start, end).
178func (d *indexData) calculateStatsForFileRange(start, end uint32) RepoStats {
179 if start >= end {
180 // An empty shard for an empty repository.
181 return RepoStats{
182 Shards: 1,
183 }
184 }
185
186 bytesContent := d.boundaries[end] - d.boundaries[start]
187 bytesFN := d.fileNameIndex[end] - d.fileNameIndex[start]
188 count, defaultCount, otherCount := d.calculateNewLinesStats(start, end)
189
190 // CR keegan for stefan: I think we may want to restructure RepoListEntry so
191 // that we don't change anything, except we have
192 // []Repository. Alternatively, things we can divide up we do (like
193 // here). Right now I don't like that these numbers are not true, especially
194 // after aggregation. For now I will move forward with this until we can
195 // chat more.
196 return RepoStats{
197 ContentBytes: int64(bytesContent) + int64(bytesFN),
198 Documents: int(end - start),
199 // CR keegan for stefan: our shard count is going to go out of whack,
200 // since we will aggregate these. So we will report more shards than are
201 // present on disk. What should we do?
202 Shards: 1,
203
204 // Sourcegraph specific
205 NewLinesCount: count,
206 DefaultBranchNewLinesCount: defaultCount,
207 OtherBranchesNewLinesCount: otherCount,
208 }
209}
210
211func (d *indexData) calculateStats() error {
212 d.repoListEntry = make([]RepoListEntry, 0, len(d.repoMetaData))
213 var start, end uint32
214
215 for repoID, md := range d.repoMetaData {
216 // determine the file range for repo i
217 for end < uint32(len(d.repos)) && d.repos[end] == uint16(repoID) {
218 end++
219 }
220 if start < end && d.repos[start] != uint16(repoID) {
221 return fmt.Errorf("shard documents out of order with respect to repositories: expected document %d to be part of repo %d", start, repoID)
222 }
223
224 d.repoListEntry = append(d.repoListEntry, RepoListEntry{
225 Repository: md,
226 IndexMetadata: d.metaData,
227 Stats: d.calculateStatsForFileRange(start, end),
228 })
229 start = end
230 }
231
232 // All repos in a compound shard share memoryUse. So we average out the
233 // memoryUse per shard in our reporting. This has the benefit that when you
234 // aggregate the IndexBytes you get back the actual memoryUse.
235 //
236 // TODO take into account tombstones for aggregation. Even better, adjust
237 // API to be shard centric not repo centric.
238 if len(d.repoListEntry) > 0 {
239 indexBytes := d.memoryUse()
240 indexBytesChunk := indexBytes / len(d.repoListEntry)
241 for i := range d.repoListEntry {
242 d.repoListEntry[i].Stats.IndexBytes = int64(indexBytesChunk)
243 indexBytes -= indexBytesChunk
244 }
245 d.repoListEntry[0].Stats.IndexBytes += int64(indexBytes)
246 }
247
248 return nil
249}
250
251// calculateNewLinesStats computes some Sourcegraph specific statistics for files
252// in the range [start, end). These are not as efficient to calculate as the
253// normal statistics. We experimentally measured about a 10% slower shard load
254// time. However, we find these values very useful to track and computing them
255// outside of load time introduces a lot of complexity.
256func (d *indexData) calculateNewLinesStats(start, end uint32) (count, defaultCount, otherCount uint64) {
257 for i := start; i < end; i++ {
258 // branchMask is a bitmask of the branches for a document. Zoekt by
259 // convention represents the default branch as the lowest bit.
260 branchMask := d.fileBranchMasks[i]
261 isDefault := (branchMask & 1) == 1
262 others := uint64(bits.OnesCount64(branchMask >> 1))
263
264 // this is readNewlines but only reading the size of each section which
265 // corresponds to the number of newlines.
266 sec := simpleSection{
267 off: d.newlinesStart + d.newlinesIndex[i],
268 sz: d.newlinesIndex[i+1] - d.newlinesIndex[i],
269 }
270 // We are only reading the first varint which is the size. So we don't
271 // need to read more than MaxVarintLen64 bytes.
272 if sec.sz > binary.MaxVarintLen64 {
273 sec.sz = binary.MaxVarintLen64
274 }
275 blob, err := d.readSectionBlob(sec)
276 if err != nil {
277 log.Printf("error reading newline index for document %d on shard %s: %v", i, d.file.Name(), err)
278 continue
279 }
280 sz, _ := binary.Uvarint(blob)
281
282 count += sz
283 if isDefault {
284 defaultCount += sz
285 }
286 otherCount += (others * sz)
287 }
288
289 return
290}
291
292func (d *indexData) String() string {
293 return fmt.Sprintf("shard(%s)", d.file.Name())
294}
295
296// calculates an approximate size of indexData in memory in bytes.
297func (d *indexData) memoryUse() int {
298 sz := 0
299 for _, a := range [][]uint32{
300 d.newlinesIndex, d.docSectionsIndex,
301 d.boundaries, d.fileNameIndex,
302 d.fileEndRunes, d.fileNameEndRunes,
303 d.fileEndSymbol, d.symbols.symKindIndex,
304 d.subRepos,
305 } {
306 sz += 4 * len(a)
307 }
308 sz += d.runeOffsets.sizeBytes()
309 sz += d.fileNameRuneOffsets.sizeBytes()
310 sz += len(d.languages)
311 sz += len(d.checksums)
312 sz += 2 * len(d.repos)
313 if len(d.ranks) > 0 {
314 sz += 8 * len(d.ranks) * len(d.ranks[0])
315 }
316 sz += 8 * len(d.runeDocSections)
317 sz += 8 * len(d.fileBranchMasks)
318 sz += d.contentNgrams.SizeBytes()
319 sz += d.fileNameNgrams.SizeBytes()
320 return sz
321}
322
323// findSelectiveNgrams returns two ngrams to pass to the distance iterator, chosen to
324// produce a small file intersection. It finds the two lowest frequency ngrams, but avoids
325// overlapping trigrams to keep their intersection as small as possible.
326//
327// Invariant: first will always have a smaller index than last.
328func findSelectiveNgrams(ngramOffs []runeNgramOff, indexMap []int, frequencies []uint32) (first, last runeNgramOff) {
329 first, last = minFrequencyNgramOffsets(ngramOffs, frequencies)
330
331 // If the trigrams are overlapping, then try to shift one to reduce overlap.
332 // This is guaranteed to produce a smaller intersection.
333 if last.index-first.index < ngramSize {
334 newFirstIndex := max(last.index-ngramSize, 0)
335 if newFirstIndex != first.index {
336 first = ngramOffs[indexMap[newFirstIndex]]
337 }
338
339 newLastIndex := min(first.index+ngramSize, len(ngramOffs)-1)
340 if newLastIndex != last.index {
341 last = ngramOffs[indexMap[newLastIndex]]
342 }
343 }
344 return
345}
346
347const maxUInt32 = 0xffffffff
348
349func minFrequencyNgramOffsets(ngramOffs []runeNgramOff, frequencies []uint32) (first, last runeNgramOff) {
350 // Find the two lowest frequency ngrams.
351 idx0, idx1 := 0, 0
352 min0, min1 := uint32(maxUInt32), uint32(maxUInt32)
353 for i, x := range frequencies {
354 if x <= min0 {
355 idx0, idx1 = i, idx0
356 min0, min1 = x, min0
357 } else if x <= min1 {
358 idx1 = i
359 min1 = x
360 }
361 }
362
363 first = ngramOffs[idx0]
364 last = ngramOffs[idx1]
365
366 // Ensure first appears before last as a helpful invariant.
367 if first.index > last.index {
368 last, first = first, last
369 }
370 return
371}
372
373func (data *indexData) ngrams(filename bool) btreeIndex {
374 if filename {
375 return data.fileNameNgrams
376 }
377 return data.contentNgrams
378}
379
380type ngramIterationResults struct {
381 matchIterator
382
383 caseSensitive bool
384 fileName bool
385 substrBytes []byte
386 substrLowered []byte
387}
388
389func (r *ngramIterationResults) String() string {
390 return fmt.Sprintf("wrapper(%v)", r.matchIterator)
391}
392
393func (r *ngramIterationResults) candidates() []*candidateMatch {
394 cs := r.matchIterator.candidates()
395 for _, c := range cs {
396 c.caseSensitive = r.caseSensitive
397 c.fileName = r.fileName
398 c.substrBytes = r.substrBytes
399 c.substrLowered = r.substrLowered
400 }
401 return cs
402}
403
404func (d *indexData) iterateNgrams(query *query.Substring) (*ngramIterationResults, error) {
405 str := query.Pattern
406
407 // Find the 2 least common ngrams from the string.
408 ngramOffs := splitNGrams([]byte(str))
409
410 // protect against accidental searching of empty strings
411 if len(ngramOffs) == 0 {
412 return nil, errors.New("iterateNgrams needs non empty string")
413 }
414
415 // PERF: Sort to increase the chances adjacent checks are in the same btree
416 // bucket (which can cause disk IO).
417 slices.SortFunc(ngramOffs, runeNgramOff.Compare)
418 frequencies := make([]uint32, 0, len(ngramOffs))
419 indexMap := make([]int, len(ngramOffs))
420 ngramLookups := 0
421 ngrams := d.ngrams(query.FileName)
422 for i, o := range ngramOffs {
423 var freq uint32
424 if query.CaseSensitive {
425 freq = ngrams.Get(o.ngram).sz
426 ngramLookups++
427 } else {
428 for _, v := range generateCaseNgrams(o.ngram) {
429 freq += ngrams.Get(v).sz
430 ngramLookups++
431 }
432 }
433
434 if freq == 0 {
435 return &ngramIterationResults{
436 matchIterator: &noMatchTree{
437 Why: "freq=0",
438 Stats: Stats{
439 NgramLookups: ngramLookups,
440 },
441 },
442 }, nil
443 }
444
445 frequencies = append(frequencies, freq)
446 indexMap[o.index] = i
447 }
448
449 first, last := findSelectiveNgrams(ngramOffs, indexMap, frequencies)
450
451 iter := &ngramDocIterator{
452 leftPad: uint32(first.index),
453 rightPad: uint32(utf8.RuneCountInString(str) - first.index),
454 ngramLookups: ngramLookups,
455 }
456 if query.FileName {
457 iter.ends = d.fileNameEndRunes
458 } else {
459 iter.ends = d.fileEndRunes
460 }
461
462 if first != last {
463 runeDist := uint32(last.index - first.index)
464 i, err := d.newDistanceTrigramIter(first.ngram, last.ngram, runeDist, query.CaseSensitive, query.FileName)
465 if err != nil {
466 return nil, err
467 }
468
469 iter.iter = i
470 } else {
471 hitIter, err := d.trigramHitIterator(last.ngram, query.CaseSensitive, query.FileName)
472 if err != nil {
473 return nil, err
474 }
475 iter.iter = hitIter
476 }
477
478 patBytes := []byte(query.Pattern)
479 lowerPatBytes := toLower(patBytes)
480
481 return &ngramIterationResults{
482 matchIterator: iter,
483 caseSensitive: query.CaseSensitive,
484 fileName: query.FileName,
485 substrBytes: patBytes,
486 substrLowered: lowerPatBytes,
487 }, nil
488}
489
490func (d *indexData) fileName(i uint32) []byte {
491 return d.fileNameContent[d.fileNameIndex[i]:d.fileNameIndex[i+1]]
492}
493
494func (d *indexData) numDocs() uint32 {
495 return uint32(len(d.fileBranchMasks))
496}
497
498func (s *indexData) Close() {
499 s.file.Close()
500}
501
502const (
503 rawConfigYes = 1
504 rawConfigNo = 2
505)
506
507// encodeRawConfig encodes a rawConfig map into a uint8 mask.
508func encodeRawConfig(rawConfig map[string]string) uint8 {
509 var encoded uint8
510 for i, f := range []string{"public", "fork", "archived"} {
511 var e uint8
512 v, ok := rawConfig[f]
513 if ok && v == "1" {
514 e |= rawConfigYes
515 } else {
516 e |= rawConfigNo
517 }
518 encoded = encoded | e<<(2*i)
519 }
520 return encoded
521}