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 "bytes"
19 "encoding/binary"
20 "fmt"
21 "hash/crc64"
22 "html/template"
23 "log"
24 "os"
25 "path/filepath"
26 "sort"
27 "time"
28 "unicode/utf8"
29
30 "github.com/go-enry/go-enry/v2"
31)
32
33var _ = log.Println
34
35const ngramSize = 3
36
37type searchableString struct {
38 data []byte
39}
40
41// Filled by the linker
42var Version string
43
44func HostnameBestEffort() string {
45 if h := os.Getenv("NODE_NAME"); h != "" {
46 return h
47 }
48 if h := os.Getenv("HOSTNAME"); h != "" {
49 return h
50 }
51 hostname, _ := os.Hostname()
52 return hostname
53}
54
55// Store character (unicode codepoint) offset (in bytes) this often.
56const runeOffsetFrequency = 100
57
58type postingsBuilder struct {
59 postings map[ngram][]byte
60 lastOffsets map[ngram]uint32
61
62 // To support UTF-8 searching, we must map back runes to byte
63 // offsets. As a first attempt, we sample regularly. The
64 // precise offset can be found by walking from the recorded
65 // offset to the desired rune.
66 runeOffsets []uint32
67 runeCount uint32
68
69 isPlainASCII bool
70
71 endRunes []uint32
72 endByte uint32
73}
74
75func newPostingsBuilder() *postingsBuilder {
76 return &postingsBuilder{
77 postings: map[ngram][]byte{},
78 lastOffsets: map[ngram]uint32{},
79 isPlainASCII: true,
80 }
81}
82
83// Store trigram offsets for the given UTF-8 data. The
84// DocumentSections must correspond to rune boundaries in the UTF-8
85// data.
86func (s *postingsBuilder) newSearchableString(data []byte, byteSections []DocumentSection) (*searchableString, []DocumentSection, error) {
87 dest := searchableString{
88 data: data,
89 }
90 var buf [8]byte
91 var runeGram [3]rune
92
93 var runeIndex uint32
94 byteCount := 0
95 dataSz := uint32(len(data))
96
97 byteSectionBoundaries := make([]uint32, 0, 2*len(byteSections))
98 for _, s := range byteSections {
99 byteSectionBoundaries = append(byteSectionBoundaries, s.Start, s.End)
100 }
101 var runeSectionBoundaries []uint32
102
103 endRune := s.runeCount
104 for ; len(data) > 0; runeIndex++ {
105 c, sz := utf8.DecodeRune(data)
106 if sz > 1 {
107 s.isPlainASCII = false
108 }
109 data = data[sz:]
110
111 runeGram[0], runeGram[1], runeGram[2] = runeGram[1], runeGram[2], c
112
113 if idx := s.runeCount + runeIndex; idx%runeOffsetFrequency == 0 {
114 s.runeOffsets = append(s.runeOffsets, s.endByte+uint32(byteCount))
115 }
116 for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] == uint32(byteCount) {
117 runeSectionBoundaries = append(runeSectionBoundaries,
118 endRune+uint32(runeIndex))
119 byteSectionBoundaries = byteSectionBoundaries[1:]
120 }
121
122 byteCount += sz
123
124 if runeIndex < 2 {
125 continue
126 }
127
128 ng := runesToNGram(runeGram)
129 lastOff := s.lastOffsets[ng]
130 newOff := endRune + uint32(runeIndex) - 2
131
132 m := binary.PutUvarint(buf[:], uint64(newOff-lastOff))
133 s.postings[ng] = append(s.postings[ng], buf[:m]...)
134 s.lastOffsets[ng] = newOff
135 }
136 s.runeCount += runeIndex
137
138 for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] < uint32(byteCount) {
139 return nil, nil, fmt.Errorf("no rune for section boundary at byte %d", byteSectionBoundaries[0])
140 }
141
142 // Handle symbol definition that ends at file end. This can
143 // happen for labels at the end of .bat files.
144
145 for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] == uint32(byteCount) {
146 runeSectionBoundaries = append(runeSectionBoundaries,
147 endRune+runeIndex)
148 byteSectionBoundaries = byteSectionBoundaries[1:]
149 }
150 runeSecs := make([]DocumentSection, 0, len(byteSections))
151 for i := 0; i < len(runeSectionBoundaries); i += 2 {
152 runeSecs = append(runeSecs, DocumentSection{
153 Start: runeSectionBoundaries[i],
154 End: runeSectionBoundaries[i+1],
155 })
156 }
157
158 s.endRunes = append(s.endRunes, s.runeCount)
159 s.endByte += dataSz
160 return &dest, runeSecs, nil
161}
162
163// IndexBuilder builds a single index shard.
164type IndexBuilder struct {
165 // The version we will write to disk. Sourcegraph Specific. This is to
166 // enable feature flagging new format versions.
167 indexFormatVersion int
168 featureVersion int
169
170 contentStrings []*searchableString
171 nameStrings []*searchableString
172 docSections [][]DocumentSection
173 runeDocSections []DocumentSection
174
175 symID uint32
176 symIndex map[string]uint32
177 symKindID uint32
178 symKindIndex map[string]uint32
179 symMetaData []uint32
180
181 fileEndSymbol []uint32
182
183 checksums []byte
184
185 branchMasks []uint64
186 subRepos []uint32
187
188 // docID => repoID
189 repos []uint16
190
191 // Experimental: docID => rank vec
192 ranks [][]float64
193
194 contentPostings *postingsBuilder
195 namePostings *postingsBuilder
196
197 // root repositories
198 repoList []Repository
199
200 // name to index.
201 subRepoIndices []map[string]uint32
202
203 // language => language code
204 languageMap map[string]uint16
205
206 // language codes, uint16 encoded as little-endian
207 languages []uint8
208
209 // IndexTime will be used as the time if non-zero. Otherwise
210 // time.Now(). This is useful for doing reproducible builds in tests.
211 IndexTime time.Time
212
213 // a sortable 20 chars long id.
214 ID string
215}
216
217func (d *Repository) verify() error {
218 for _, t := range []string{d.FileURLTemplate, d.LineFragmentTemplate, d.CommitURLTemplate} {
219 if _, err := template.New("").Parse(t); err != nil {
220 return err
221 }
222 }
223 return nil
224}
225
226// ContentSize returns the number of content bytes so far ingested.
227func (b *IndexBuilder) ContentSize() uint32 {
228 // Add the name too so we don't skip building index if we have
229 // lots of empty files.
230 return b.contentPostings.endByte + b.namePostings.endByte
231}
232
233// NewIndexBuilder creates a fresh IndexBuilder. The passed in
234// Repository contains repo metadata, and may be set to nil.
235func NewIndexBuilder(r *Repository) (*IndexBuilder, error) {
236 b := newIndexBuilder()
237
238 if r == nil {
239 r = &Repository{}
240 }
241 if err := b.setRepository(r); err != nil {
242 return nil, err
243 }
244 return b, nil
245}
246
247func newIndexBuilder() *IndexBuilder {
248 return &IndexBuilder{
249 indexFormatVersion: IndexFormatVersion,
250 featureVersion: FeatureVersion,
251
252 contentPostings: newPostingsBuilder(),
253 namePostings: newPostingsBuilder(),
254 fileEndSymbol: []uint32{0},
255 symIndex: make(map[string]uint32),
256 symKindIndex: make(map[string]uint32),
257 languageMap: make(map[string]uint16),
258 }
259}
260
261func (b *IndexBuilder) setRepository(desc *Repository) error {
262 if err := desc.verify(); err != nil {
263 return err
264 }
265
266 if len(desc.Branches) > 64 {
267 return fmt.Errorf("too many branches")
268 }
269
270 repo := *desc
271
272 // copy subrepomap without root
273 repo.SubRepoMap = map[string]*Repository{}
274 for k, v := range desc.SubRepoMap {
275 if k != "" {
276 repo.SubRepoMap[k] = v
277 }
278 }
279
280 b.repoList = append(b.repoList, repo)
281
282 return b.populateSubRepoIndices()
283}
284
285type DocumentSection struct {
286 Start, End uint32
287}
288
289// Document holds a document (file) to index.
290type Document struct {
291 Name string
292 Content []byte
293 Branches []string
294 SubRepositoryPath string
295 Language string
296
297 // If set, something is wrong with the file contents, and this
298 // is the reason it wasn't indexed.
299 SkipReason string
300
301 // Document sections for symbols. Offsets should use bytes.
302 Symbols []DocumentSection
303 SymbolsMetaData []*Symbol
304
305 // Ranks is a vector of ranks for a document as provided by a DocumentRanksFile
306 // file in the git repo.
307 //
308 // Two documents can be ordered by comparing the components of their rank
309 // vectors. Bigger entries are better, as are longer vectors.
310 //
311 // This field is experimental and may change at any time without warning.
312 Ranks []float64
313}
314
315type symbolSlice struct {
316 symbols []DocumentSection
317 metaData []*Symbol
318}
319
320func (s symbolSlice) Len() int { return len(s.symbols) }
321
322func (s symbolSlice) Swap(i, j int) {
323 s.symbols[i], s.symbols[j] = s.symbols[j], s.symbols[i]
324 s.metaData[i], s.metaData[j] = s.metaData[j], s.metaData[i]
325}
326
327func (s symbolSlice) Less(i, j int) bool {
328 return s.symbols[i].Start < s.symbols[j].Start
329}
330
331// AddFile is a convenience wrapper for Add
332func (b *IndexBuilder) AddFile(name string, content []byte) error {
333 return b.Add(Document{Name: name, Content: content})
334}
335
336// CheckText returns a reason why the given contents are probably not source texts.
337func CheckText(content []byte, maxTrigramCount int) error {
338 if len(content) == 0 {
339 return nil
340 }
341
342 if len(content) < ngramSize {
343 return fmt.Errorf("file size smaller than %d", ngramSize)
344 }
345
346 // PERF: we only need to do the trigram check if the upperbound on content
347 // is greater than our threshold.
348 var trigrams map[ngram]struct{}
349 if trigramsUpperBound := len(content) - ngramSize + 1; trigramsUpperBound > maxTrigramCount {
350 trigrams = make(map[ngram]struct{}, maxTrigramCount+1)
351 }
352
353 var cur [3]rune
354 byteCount := 0
355 for len(content) > 0 {
356 if content[0] == 0 {
357 return fmt.Errorf("binary data at byte offset %d", byteCount)
358 }
359
360 r, sz := utf8.DecodeRune(content)
361 content = content[sz:]
362 byteCount += sz
363
364 cur[0], cur[1], cur[2] = cur[1], cur[2], r
365 if cur[0] == 0 {
366 // start of file.
367 continue
368 }
369
370 if trigrams != nil {
371 trigrams[runesToNGram(cur)] = struct{}{}
372 if len(trigrams) > maxTrigramCount {
373 // probably not text.
374 return fmt.Errorf("number of trigrams exceeds %d", maxTrigramCount)
375 }
376 }
377 }
378 return nil
379}
380
381func (b *IndexBuilder) populateSubRepoIndices() error {
382 if len(b.subRepoIndices) == len(b.repoList) {
383 return nil
384 }
385 if len(b.subRepoIndices) != len(b.repoList)-1 {
386 return fmt.Errorf("populateSubRepoIndices not called for a repo: %d != %d - 1", len(b.subRepoIndices), len(b.repoList))
387 }
388 repo := b.repoList[len(b.repoList)-1]
389 b.subRepoIndices = append(b.subRepoIndices, mkSubRepoIndices(repo))
390 return nil
391}
392
393func mkSubRepoIndices(repo Repository) map[string]uint32 {
394 paths := []string{""}
395 for k := range repo.SubRepoMap {
396 paths = append(paths, k)
397 }
398 sort.Strings(paths)
399 subRepoIndices := make(map[string]uint32, len(paths))
400 for i, p := range paths {
401 subRepoIndices[p] = uint32(i)
402 }
403 return subRepoIndices
404}
405
406const notIndexedMarker = "NOT-INDEXED: "
407
408func (b *IndexBuilder) symbolID(sym string) uint32 {
409 if _, ok := b.symIndex[sym]; !ok {
410 b.symIndex[sym] = b.symID
411 b.symID++
412 }
413 return b.symIndex[sym]
414}
415
416func (b *IndexBuilder) symbolKindID(t string) uint32 {
417 if _, ok := b.symKindIndex[t]; !ok {
418 b.symKindIndex[t] = b.symKindID
419 b.symKindID++
420 }
421 return b.symKindIndex[t]
422}
423
424func (b *IndexBuilder) addSymbols(symbols []*Symbol) {
425 for _, sym := range symbols {
426 b.symMetaData = append(b.symMetaData,
427 // This field was removed due to redundancy. To avoid
428 // needing to reindex, it is set to zero for now. In the
429 // future, this field will be completely removed. It
430 // will require incrementing the feature version.
431 0,
432 b.symbolKindID(sym.Kind),
433 b.symbolID(sym.Parent),
434 b.symbolKindID(sym.ParentKind))
435 }
436}
437
438func DetermineLanguageIfUnknown(doc *Document) {
439 if doc.Language == "" {
440 c := doc.Content
441 // classifier is faster on small files without losing much accuracy
442 if len(c) > 2048 {
443 c = c[:2048]
444 }
445 doc.Language = enry.GetLanguage(doc.Name, c)
446 }
447}
448
449// Add a file which only occurs in certain branches.
450func (b *IndexBuilder) Add(doc Document) error {
451 hasher := crc64.New(crc64.MakeTable(crc64.ISO))
452
453 if idx := bytes.IndexByte(doc.Content, 0); idx >= 0 {
454 doc.SkipReason = fmt.Sprintf("binary content at byte offset %d", idx)
455 doc.Language = "binary"
456 }
457
458 if doc.SkipReason != "" {
459 doc.Content = []byte(notIndexedMarker + doc.SkipReason)
460 doc.Symbols = nil
461 doc.SymbolsMetaData = nil
462 if doc.Language == "" {
463 doc.Language = "skipped"
464 }
465 }
466
467 DetermineLanguageIfUnknown(&doc)
468
469 sort.Sort(symbolSlice{doc.Symbols, doc.SymbolsMetaData})
470 var last DocumentSection
471 for i, s := range doc.Symbols {
472 if i > 0 {
473 if last.End > s.Start {
474 return fmt.Errorf("sections overlap")
475 }
476 }
477 last = s
478 }
479 if last.End > uint32(len(doc.Content)) {
480 return fmt.Errorf("section goes past end of content")
481 }
482
483 if doc.SubRepositoryPath != "" {
484 rel, err := filepath.Rel(doc.SubRepositoryPath, doc.Name)
485 if err != nil || rel == doc.Name {
486 return fmt.Errorf("path %q must start subrepo path %q", doc.Name, doc.SubRepositoryPath)
487 }
488 }
489 docStr, runeSecs, err := b.contentPostings.newSearchableString(doc.Content, doc.Symbols)
490 if err != nil {
491 return err
492 }
493 nameStr, _, err := b.namePostings.newSearchableString([]byte(doc.Name), nil)
494 if err != nil {
495 return err
496 }
497 b.addSymbols(doc.SymbolsMetaData)
498
499 repoIdx := len(b.repoList) - 1
500 subRepoIdx, ok := b.subRepoIndices[repoIdx][doc.SubRepositoryPath]
501 if !ok {
502 return fmt.Errorf("unknown subrepo path %q", doc.SubRepositoryPath)
503 }
504
505 var mask uint64
506 for _, br := range doc.Branches {
507 m := b.branchMask(br)
508 if m == 0 {
509 return fmt.Errorf("no branch found for %s", br)
510 }
511 mask |= m
512 }
513
514 if repoIdx > 1<<16 {
515 return fmt.Errorf("too many repos in shard: max is %d", 1<<16)
516 }
517
518 b.subRepos = append(b.subRepos, subRepoIdx)
519 b.repos = append(b.repos, uint16(repoIdx))
520
521 // doc.Ranks might be nil. In case we don't use offline ranking, doc.Ranks is
522 // always nil.
523 b.ranks = append(b.ranks, doc.Ranks)
524
525 hasher.Write(doc.Content)
526
527 b.contentStrings = append(b.contentStrings, docStr)
528 b.runeDocSections = append(b.runeDocSections, runeSecs...)
529
530 b.nameStrings = append(b.nameStrings, nameStr)
531 b.docSections = append(b.docSections, doc.Symbols)
532 b.fileEndSymbol = append(b.fileEndSymbol, uint32(len(b.runeDocSections)))
533 b.branchMasks = append(b.branchMasks, mask)
534 b.checksums = append(b.checksums, hasher.Sum(nil)...)
535
536 langCode, ok := b.languageMap[doc.Language]
537 if !ok {
538 if len(b.languageMap) >= 65535 {
539 return fmt.Errorf("too many languages")
540 }
541 langCode = uint16(len(b.languageMap))
542 b.languageMap[doc.Language] = langCode
543 }
544 b.languages = append(b.languages, uint8(langCode), uint8(langCode>>8))
545
546 return nil
547}
548
549func (b *IndexBuilder) branchMask(br string) uint64 {
550 for i, b := range b.repoList[len(b.repoList)-1].Branches {
551 if b.Name == br {
552 return uint64(1) << uint(i)
553 }
554 }
555 return 0
556}