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// NumFiles returns the number of files added to this builder
234func (b *IndexBuilder) NumFiles() int {
235 return len(b.contentStrings)
236}
237
238// NewIndexBuilder creates a fresh IndexBuilder. The passed in
239// Repository contains repo metadata, and may be set to nil.
240func NewIndexBuilder(r *Repository) (*IndexBuilder, error) {
241 b := newIndexBuilder()
242
243 if r == nil {
244 r = &Repository{}
245 }
246 if err := b.setRepository(r); err != nil {
247 return nil, err
248 }
249 return b, nil
250}
251
252func newIndexBuilder() *IndexBuilder {
253 return &IndexBuilder{
254 indexFormatVersion: IndexFormatVersion,
255 featureVersion: FeatureVersion,
256
257 contentPostings: newPostingsBuilder(),
258 namePostings: newPostingsBuilder(),
259 fileEndSymbol: []uint32{0},
260 symIndex: make(map[string]uint32),
261 symKindIndex: make(map[string]uint32),
262 languageMap: make(map[string]uint16),
263 }
264}
265
266func (b *IndexBuilder) setRepository(desc *Repository) error {
267 if err := desc.verify(); err != nil {
268 return err
269 }
270
271 if len(desc.Branches) > 64 {
272 return fmt.Errorf("too many branches")
273 }
274
275 repo := *desc
276
277 // copy subrepomap without root
278 repo.SubRepoMap = map[string]*Repository{}
279 for k, v := range desc.SubRepoMap {
280 if k != "" {
281 repo.SubRepoMap[k] = v
282 }
283 }
284
285 b.repoList = append(b.repoList, repo)
286
287 return b.populateSubRepoIndices()
288}
289
290type DocumentSection struct {
291 Start, End uint32
292}
293
294// Document holds a document (file) to index.
295type Document struct {
296 Name string
297 Content []byte
298 Branches []string
299 SubRepositoryPath string
300 Language string
301
302 // If set, something is wrong with the file contents, and this
303 // is the reason it wasn't indexed.
304 SkipReason string
305
306 // Document sections for symbols. Offsets should use bytes.
307 Symbols []DocumentSection
308 SymbolsMetaData []*Symbol
309
310 // Ranks is a vector of ranks for a document as provided by a DocumentRanksFile
311 // file in the git repo.
312 //
313 // Two documents can be ordered by comparing the components of their rank
314 // vectors. Bigger entries are better, as are longer vectors.
315 //
316 // This field is experimental and may change at any time without warning.
317 Ranks []float64
318}
319
320type symbolSlice struct {
321 symbols []DocumentSection
322 metaData []*Symbol
323}
324
325func (s symbolSlice) Len() int { return len(s.symbols) }
326
327func (s symbolSlice) Swap(i, j int) {
328 s.symbols[i], s.symbols[j] = s.symbols[j], s.symbols[i]
329 s.metaData[i], s.metaData[j] = s.metaData[j], s.metaData[i]
330}
331
332func (s symbolSlice) Less(i, j int) bool {
333 return s.symbols[i].Start < s.symbols[j].Start
334}
335
336// AddFile is a convenience wrapper for Add
337func (b *IndexBuilder) AddFile(name string, content []byte) error {
338 return b.Add(Document{Name: name, Content: content})
339}
340
341func (b *IndexBuilder) populateSubRepoIndices() error {
342 if len(b.subRepoIndices) == len(b.repoList) {
343 return nil
344 }
345 if len(b.subRepoIndices) != len(b.repoList)-1 {
346 return fmt.Errorf("populateSubRepoIndices not called for a repo: %d != %d - 1", len(b.subRepoIndices), len(b.repoList))
347 }
348 repo := b.repoList[len(b.repoList)-1]
349 b.subRepoIndices = append(b.subRepoIndices, mkSubRepoIndices(repo))
350 return nil
351}
352
353func mkSubRepoIndices(repo Repository) map[string]uint32 {
354 paths := []string{""}
355 for k := range repo.SubRepoMap {
356 paths = append(paths, k)
357 }
358 sort.Strings(paths)
359 subRepoIndices := make(map[string]uint32, len(paths))
360 for i, p := range paths {
361 subRepoIndices[p] = uint32(i)
362 }
363 return subRepoIndices
364}
365
366const notIndexedMarker = "NOT-INDEXED: "
367
368func (b *IndexBuilder) symbolID(sym string) uint32 {
369 if _, ok := b.symIndex[sym]; !ok {
370 b.symIndex[sym] = b.symID
371 b.symID++
372 }
373 return b.symIndex[sym]
374}
375
376func (b *IndexBuilder) symbolKindID(t string) uint32 {
377 if _, ok := b.symKindIndex[t]; !ok {
378 b.symKindIndex[t] = b.symKindID
379 b.symKindID++
380 }
381 return b.symKindIndex[t]
382}
383
384func (b *IndexBuilder) addSymbols(symbols []*Symbol) {
385 for _, sym := range symbols {
386 b.symMetaData = append(b.symMetaData,
387 // This field was removed due to redundancy. To avoid
388 // needing to reindex, it is set to zero for now. In the
389 // future, this field will be completely removed. It
390 // will require incrementing the feature version.
391 0,
392 b.symbolKindID(sym.Kind),
393 b.symbolID(sym.Parent),
394 b.symbolKindID(sym.ParentKind))
395 }
396}
397
398func DetermineLanguageIfUnknown(doc *Document) {
399 if doc.Language == "" {
400 doc.Language = enry.GetLanguage(doc.Name, doc.Content)
401 }
402}
403
404// Add a file which only occurs in certain branches.
405func (b *IndexBuilder) Add(doc Document) error {
406 hasher := crc64.New(crc64.MakeTable(crc64.ISO))
407
408 if idx := bytes.IndexByte(doc.Content, 0); idx >= 0 {
409 doc.SkipReason = fmt.Sprintf("binary content at byte offset %d", idx)
410 doc.Language = "binary"
411 }
412
413 if doc.SkipReason != "" {
414 doc.Content = []byte(notIndexedMarker + doc.SkipReason)
415 doc.Symbols = nil
416 doc.SymbolsMetaData = nil
417 if doc.Language == "" {
418 doc.Language = "skipped"
419 }
420 }
421
422 DetermineLanguageIfUnknown(&doc)
423
424 sort.Sort(symbolSlice{doc.Symbols, doc.SymbolsMetaData})
425 var last DocumentSection
426 for i, s := range doc.Symbols {
427 if i > 0 {
428 if last.End > s.Start {
429 return fmt.Errorf("sections overlap")
430 }
431 }
432 last = s
433 }
434 if last.End > uint32(len(doc.Content)) {
435 return fmt.Errorf("section goes past end of content")
436 }
437
438 if doc.SubRepositoryPath != "" {
439 rel, err := filepath.Rel(doc.SubRepositoryPath, doc.Name)
440 if err != nil || rel == doc.Name {
441 return fmt.Errorf("path %q must start subrepo path %q", doc.Name, doc.SubRepositoryPath)
442 }
443 }
444 docStr, runeSecs, err := b.contentPostings.newSearchableString(doc.Content, doc.Symbols)
445 if err != nil {
446 return err
447 }
448 nameStr, _, err := b.namePostings.newSearchableString([]byte(doc.Name), nil)
449 if err != nil {
450 return err
451 }
452 b.addSymbols(doc.SymbolsMetaData)
453
454 repoIdx := len(b.repoList) - 1
455 subRepoIdx, ok := b.subRepoIndices[repoIdx][doc.SubRepositoryPath]
456 if !ok {
457 return fmt.Errorf("unknown subrepo path %q", doc.SubRepositoryPath)
458 }
459
460 var mask uint64
461 for _, br := range doc.Branches {
462 m := b.branchMask(br)
463 if m == 0 {
464 return fmt.Errorf("no branch found for %s", br)
465 }
466 mask |= m
467 }
468
469 if repoIdx > 1<<16 {
470 return fmt.Errorf("too many repos in shard: max is %d", 1<<16)
471 }
472
473 b.subRepos = append(b.subRepos, subRepoIdx)
474 b.repos = append(b.repos, uint16(repoIdx))
475
476 // doc.Ranks might be nil. In case we don't use offline ranking, doc.Ranks is
477 // always nil.
478 b.ranks = append(b.ranks, doc.Ranks)
479
480 hasher.Write(doc.Content)
481
482 b.contentStrings = append(b.contentStrings, docStr)
483 b.runeDocSections = append(b.runeDocSections, runeSecs...)
484
485 b.nameStrings = append(b.nameStrings, nameStr)
486 b.docSections = append(b.docSections, doc.Symbols)
487 b.fileEndSymbol = append(b.fileEndSymbol, uint32(len(b.runeDocSections)))
488 b.branchMasks = append(b.branchMasks, mask)
489 b.checksums = append(b.checksums, hasher.Sum(nil)...)
490
491 langCode, ok := b.languageMap[doc.Language]
492 if !ok {
493 if len(b.languageMap) >= 65535 {
494 return fmt.Errorf("too many languages")
495 }
496 langCode = uint16(len(b.languageMap))
497 b.languageMap[doc.Language] = langCode
498 }
499 b.languages = append(b.languages, uint8(langCode), uint8(langCode>>8))
500
501 return nil
502}
503
504func (b *IndexBuilder) branchMask(br string) uint64 {
505 for i, b := range b.repoList[len(b.repoList)-1].Branches {
506 if b.Name == br {
507 return uint64(1) << uint(i)
508 }
509 }
510 return 0
511}
512
513type DocChecker struct {
514 // A map to count the unique trigrams in a doc. Reused across docs to cut down on allocations.
515 trigrams map[ngram]struct{}
516}
517
518// Check returns a reason why the given contents are probably not source texts.
519func (t *DocChecker) Check(content []byte, maxTrigramCount int, allowLargeFile bool) error {
520 if len(content) == 0 {
521 return nil
522 }
523
524 if len(content) < ngramSize {
525 return fmt.Errorf("file size smaller than %d", ngramSize)
526 }
527
528 if index := bytes.IndexByte(content, 0); index > 0 {
529 return fmt.Errorf("binary data at byte offset %d", index)
530 }
531
532 // PERF: we only need to do the trigram check if the upperbound on content is greater than
533 // our threshold. Also skip the trigram check if the file is explicitly marked as allowed.
534 if trigramsUpperBound := len(content) - ngramSize + 1; trigramsUpperBound <= maxTrigramCount || allowLargeFile {
535 return nil
536 }
537
538 var cur [3]rune
539 byteCount := 0
540 t.clearTrigrams(maxTrigramCount)
541
542 for len(content) > 0 {
543 r, sz := utf8.DecodeRune(content)
544 content = content[sz:]
545 byteCount += sz
546
547 cur[0], cur[1], cur[2] = cur[1], cur[2], r
548 if cur[0] == 0 {
549 // start of file.
550 continue
551 }
552
553 t.trigrams[runesToNGram(cur)] = struct{}{}
554 if len(t.trigrams) > maxTrigramCount {
555 // probably not text.
556 return fmt.Errorf("number of trigrams exceeds %d", maxTrigramCount)
557 }
558 }
559 return nil
560}
561
562func (t *DocChecker) clearTrigrams(maxTrigramCount int) {
563 if t.trigrams == nil {
564 t.trigrams = make(map[ngram]struct{}, maxTrigramCount)
565 }
566 for key := range t.trigrams {
567 delete(t.trigrams, key)
568 }
569}