// Copyright 2016 Google Inc. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package index import ( "bytes" "encoding/binary" "fmt" "hash/crc64" "log" "net/url" "os" "path/filepath" "slices" "sort" "strings" "text/template" "time" "unicode/utf8" "github.com/sourcegraph/zoekt" "github.com/sourcegraph/zoekt/languages" ) var _ = log.Println const ngramSize = 3 type searchableString struct { data []byte } // Filled by the linker var Version string func HostnameBestEffort() string { if h := os.Getenv("NODE_NAME"); h != "" { return h } if h := os.Getenv("HOSTNAME"); h != "" { return h } hostname, _ := os.Hostname() return hostname } // Store character (unicode codepoint) offset (in bytes) this often. const runeOffsetFrequency = 100 // postingList holds the varint-encoded delta data and last offset for a // single ngram. Stored by pointer in the asciiPostings array or the // postings map so appending to data does not require rewriting the // map entry or array slot. type postingList struct { data []byte lastOff uint32 } // asciiNgramBits is the number of bits needed to index all ASCII trigrams. // ASCII runes are 0-127 (7 bits), so 3 runes = 21 bits = 2M entries. const asciiNgramBits = 21 // asciiNgramIndex packs three ASCII bytes into a 21-bit array index. func asciiNgramIndex(a, b, c byte) uint32 { return uint32(a)<<14 | uint32(b)<<7 | uint32(c) } // asciiIndexToNgram converts a 21-bit ASCII array index back to the // canonical ngram encoding (rune[0]<<42 | rune[1]<<21 | rune[2]). func asciiIndexToNgram(idx uint32) ngram { r0 := uint64(idx >> 14) r1 := uint64((idx >> 7) & 0x7f) r2 := uint64(idx & 0x7f) return ngram(r0<<42 | r1<<21 | r2) } type postingsBuilder struct { // ASCII trigrams use direct-indexed array (zero hash/probe cost). // Non-ASCII trigrams fall back to the map. asciiPostings [1 << asciiNgramBits]*postingList postings map[ngram]*postingList // asciiPopulated tracks which indices in asciiPostings are non-nil, // so reset() and writePostings iterate only populated slots — O(n) // where n is unique ASCII trigrams (~275K) instead of O(2M). asciiPopulated []uint32 // To support UTF-8 searching, we must map back runes to byte // offsets. As a first attempt, we sample regularly. The // precise offset can be found by walking from the recorded // offset to the desired rune. runeOffsets []uint32 runeCount uint32 isPlainASCII bool endRunes []uint32 endByte uint32 } // Initial capacity for each posting list's byte slice. On the // kubernetes corpus (282K unique trigrams), the median posting list is // 10 bytes and 78% are under 64 bytes (power-law distribution). // Pre-allocating 64 covers the majority without the 244 MB waste that // a mean-based value (1024) would cause. const initialPostingCap = 64 // estimateNgrams returns a pre-size hint for the non-ASCII postings map, // derived from the maximum shard content size. Intentionally over-estimates // (the map only holds non-ASCII trigrams) to avoid rehashing. func estimateNgrams(shardMaxBytes int) int { n := shardMaxBytes / 600 if n < 1024 { n = 1024 } return n } func newPostingsBuilder(shardMaxBytes int) *postingsBuilder { return &postingsBuilder{ postings: make(map[ngram]*postingList, estimateNgrams(shardMaxBytes)), isPlainASCII: true, } } // reset clears the builder for reuse. All postingList allocations // (backing arrays, map entries, ASCII array slots) are retained so the // next shard build avoids re-allocating them. // Uses asciiPopulated to reset only populated slots — O(populated) // instead of O(2M). Slots are kept non-nil with data truncated to // len 0; the hot path uses len(pl.data)==0 to re-record them in // asciiPopulated for the next shard. func (s *postingsBuilder) reset() { for _, idx := range s.asciiPopulated { pl := s.asciiPostings[idx] pl.data = pl.data[:0] pl.lastOff = 0 } s.asciiPopulated = s.asciiPopulated[:0] for _, pl := range s.postings { pl.data = pl.data[:0] pl.lastOff = 0 } s.runeOffsets = s.runeOffsets[:0] s.runeCount = 0 s.isPlainASCII = true s.endRunes = s.endRunes[:0] s.endByte = 0 } // Store trigram offsets for the given UTF-8 data. The // DocumentSections must correspond to rune boundaries in the UTF-8 // data. func (s *postingsBuilder) newSearchableString(data []byte, byteSections []DocumentSection) (*searchableString, []DocumentSection, error) { dest := searchableString{ data: data, } var buf [8]byte var runeGram [3]rune var runeIndex uint32 byteCount := 0 dataSz := uint32(len(data)) byteSectionBoundaries := make([]uint32, 0, 2*len(byteSections)) for _, s := range byteSections { byteSectionBoundaries = append(byteSectionBoundaries, s.Start, s.End) } var runeSectionBoundaries []uint32 endRune := s.runeCount for ; len(data) > 0; runeIndex++ { // ASCII fast path: avoid utf8.DecodeRune call overhead. // For source code, 95-99% of bytes are ASCII. var c rune sz := 1 if data[0] < utf8.RuneSelf { c = rune(data[0]) } else { c, sz = utf8.DecodeRune(data) s.isPlainASCII = false } data = data[sz:] runeGram[0], runeGram[1], runeGram[2] = runeGram[1], runeGram[2], c if idx := s.runeCount + runeIndex; idx%runeOffsetFrequency == 0 { s.runeOffsets = append(s.runeOffsets, s.endByte+uint32(byteCount)) } for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] == uint32(byteCount) { runeSectionBoundaries = append(runeSectionBoundaries, endRune+uint32(runeIndex)) byteSectionBoundaries = byteSectionBoundaries[1:] } byteCount += sz if runeIndex < 2 { continue } newOff := endRune + uint32(runeIndex) - 2 // ASCII trigrams use direct-indexed array (no hash/probe). var pl *postingList if runeGram[0] < utf8.RuneSelf && runeGram[1] < utf8.RuneSelf && runeGram[2] < utf8.RuneSelf { idx := asciiNgramIndex(byte(runeGram[0]), byte(runeGram[1]), byte(runeGram[2])) pl = s.asciiPostings[idx] if pl == nil { pl = &postingList{data: make([]byte, 0, initialPostingCap)} s.asciiPostings[idx] = pl s.asciiPopulated = append(s.asciiPopulated, idx) } else if len(pl.data) == 0 { // Retained from a previous shard (pool reuse) — re-record // in asciiPopulated for this shard's writePostings. s.asciiPopulated = append(s.asciiPopulated, idx) } } else { ng := runesToNGram(runeGram) pl = s.postings[ng] if pl == nil { pl = &postingList{data: make([]byte, 0, initialPostingCap)} s.postings[ng] = pl } } delta := uint64(newOff - pl.lastOff) if delta < 0x80 { // Single-byte varint fast path: ~80% of deltas are < 128. // append(slice, byte) is cheaper than append(slice, slice...). pl.data = append(pl.data, byte(delta)) } else { m := binary.PutUvarint(buf[:], delta) pl.data = append(pl.data, buf[:m]...) } pl.lastOff = newOff } s.runeCount += runeIndex for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] < uint32(byteCount) { return nil, nil, fmt.Errorf("no rune for section boundary at byte %d", byteSectionBoundaries[0]) } // Handle symbol definition that ends at file end. This can // happen for labels at the end of .bat files. for len(byteSectionBoundaries) > 0 && byteSectionBoundaries[0] == uint32(byteCount) { runeSectionBoundaries = append(runeSectionBoundaries, endRune+runeIndex) byteSectionBoundaries = byteSectionBoundaries[1:] } runeSecs := make([]DocumentSection, 0, len(byteSections)) for i := 0; i < len(runeSectionBoundaries); i += 2 { runeSecs = append(runeSecs, DocumentSection{ Start: runeSectionBoundaries[i], End: runeSectionBoundaries[i+1], }) } s.endRunes = append(s.endRunes, s.runeCount) s.endByte += dataSz return &dest, runeSecs, nil } // ShardBuilder builds a single index shard. type ShardBuilder struct { // The version we will write to disk. Sourcegraph Specific. This is to // enable feature flagging new format versions. indexFormatVersion int featureVersion int contentStrings []*searchableString nameStrings []*searchableString docSections [][]DocumentSection runeDocSections []DocumentSection symID uint32 symIndex map[string]uint32 symKindID uint32 symKindIndex map[string]uint32 symMetaData []uint32 fileEndSymbol []uint32 checksums []byte branchMasks []uint64 subRepos []uint32 // docID => repoID repos []uint16 contentPostings *postingsBuilder namePostings *postingsBuilder // root repositories repoList []zoekt.Repository // name to index. subRepoIndices []map[string]uint32 // language => language code languageMap map[string]uint16 // language codes, uint16 encoded as little-endian languages []uint8 categories []byte // IndexTime will be used as the time if non-zero. Otherwise // time.Now(). This is useful for doing reproducible builds in tests. IndexTime time.Time // a sortable 20 chars long id. ID string } func verify(repo *zoekt.Repository) error { for _, t := range []string{repo.FileURLTemplate, repo.LineFragmentTemplate, repo.CommitURLTemplate} { if _, err := ParseTemplate(t); err != nil { return err } } return nil } func urlJoinPath(base string, elem ...string) string { // golangs html/template always escapes "+" appearing in an HTML attribute // [1]. We may even want to treat more characters, differently but this // atleast makes it possible to visit URLs like [2]. // // We only do this to elem since base will normally be a hardcoded string. // // [1]: https://sourcegraph.com/github.com/golang/go@go1.23.2/-/blob/src/html/template/html.go?L71-80 // [2]: https://github.com/apple/swift-system/blob/main/Sources/System/Util+StringArray.swift elem = slices.Clone(elem) // copy to mutate for i := range elem { elem[i] = strings.ReplaceAll(elem[i], "+", "%2B") } u, err := url.JoinPath(base, elem...) if err != nil { return "#!error: " + err.Error() } return u } // ParseTemplate will parse the templates for FileURLTemplate, // LineFragmentTemplate and CommitURLTemplate. // // It makes available the extra function UrlJoinPath. func ParseTemplate(text string) (*template.Template, error) { return template.New("").Funcs(template.FuncMap{ "URLJoinPath": urlJoinPath, }).Parse(text) } // ContentSize returns the number of content bytes so far ingested. func (b *ShardBuilder) ContentSize() uint32 { // Add the name too so we don't skip building index if we have // lots of empty files. return b.contentPostings.endByte + b.namePostings.endByte } // NumFiles returns the number of files added to this builder func (b *ShardBuilder) NumFiles() int { return len(b.contentStrings) } // NewShardBuilder creates a fresh ShardBuilder. The passed in // Repository contains repo metadata, and may be set to nil. func NewShardBuilder(r *zoekt.Repository) (*ShardBuilder, error) { b := newShardBuilder(0) if r == nil { r = &zoekt.Repository{} } if err := b.setRepository(r); err != nil { return nil, err } return b, nil } const defaultShardMax = 100 << 20 // 100 MB, matches Options.ShardMax default // newShardBuilder creates a ShardBuilder with fresh postingsBuilders. // shardMax is the maximum shard content size in bytes (0 uses defaultShardMax). func newShardBuilder(shardMax int) *ShardBuilder { if shardMax <= 0 { shardMax = defaultShardMax } return newShardBuilderWithPostings( newPostingsBuilder(shardMax), newPostingsBuilder(shardMax), ) } func newShardBuilderWithPostings(content, name *postingsBuilder) *ShardBuilder { return &ShardBuilder{ indexFormatVersion: IndexFormatVersion, featureVersion: FeatureVersion, contentPostings: content, namePostings: name, fileEndSymbol: []uint32{0}, symIndex: make(map[string]uint32), symKindIndex: make(map[string]uint32), languageMap: make(map[string]uint16), } } func (b *ShardBuilder) setRepository(desc *zoekt.Repository) error { if err := verify(desc); err != nil { return err } if len(desc.Branches) > 64 { return fmt.Errorf("too many branches") } repo := *desc // copy subrepomap without root repo.SubRepoMap = map[string]*zoekt.Repository{} for k, v := range desc.SubRepoMap { if k != "" { repo.SubRepoMap[k] = v } } b.repoList = append(b.repoList, repo) return b.populateSubRepoIndices() } type symbolSlice struct { symbols []DocumentSection metaData []*zoekt.Symbol } func (s symbolSlice) Len() int { return len(s.symbols) } func (s symbolSlice) Swap(i, j int) { s.symbols[i], s.symbols[j] = s.symbols[j], s.symbols[i] s.metaData[i], s.metaData[j] = s.metaData[j], s.metaData[i] } func (s symbolSlice) Less(i, j int) bool { return s.symbols[i].Start < s.symbols[j].Start } // AddFile is a convenience wrapper for Add func (b *ShardBuilder) AddFile(name string, content []byte) error { return b.Add(Document{Name: name, Content: content}) } func (b *ShardBuilder) populateSubRepoIndices() error { if len(b.subRepoIndices) == len(b.repoList) { return nil } if len(b.subRepoIndices) != len(b.repoList)-1 { return fmt.Errorf("populateSubRepoIndices not called for a repo: %d != %d - 1", len(b.subRepoIndices), len(b.repoList)) } repo := b.repoList[len(b.repoList)-1] b.subRepoIndices = append(b.subRepoIndices, mkSubRepoIndices(repo)) return nil } func mkSubRepoIndices(repo zoekt.Repository) map[string]uint32 { paths := []string{""} for k := range repo.SubRepoMap { paths = append(paths, k) } sort.Strings(paths) subRepoIndices := make(map[string]uint32, len(paths)) for i, p := range paths { subRepoIndices[p] = uint32(i) } return subRepoIndices } const notIndexedMarker = "NOT-INDEXED: " func (b *ShardBuilder) symbolID(sym string) uint32 { if _, ok := b.symIndex[sym]; !ok { b.symIndex[sym] = b.symID b.symID++ } return b.symIndex[sym] } func (b *ShardBuilder) symbolKindID(t string) uint32 { if _, ok := b.symKindIndex[t]; !ok { b.symKindIndex[t] = b.symKindID b.symKindID++ } return b.symKindIndex[t] } func (b *ShardBuilder) addSymbols(symbols []*zoekt.Symbol) { for _, sym := range symbols { b.symMetaData = append(b.symMetaData, // This field was removed due to redundancy. To avoid // needing to reindex, it is set to zero for now. In the // future, this field will be completely removed. It // will require incrementing the feature version. 0, b.symbolKindID(sym.Kind), b.symbolID(sym.Parent), b.symbolKindID(sym.ParentKind)) } } func DetermineLanguageIfUnknown(doc *Document) { if doc.Language != "" { return } // If this document has been skipped (doc.SkipReason != SkipReasonNone), it's // likely very large, or it's a non-code file like binary. In this case, we just // guess the language based on the file name to avoid examining the contents. // Note: passing nil content is allowed by the go-enry contract (the underlying // library we use here). var content []byte if doc.SkipReason == SkipReasonNone { content = doc.Content } langs := languages.GetLanguagesFromContent(doc.Name, content) if len(langs) > 0 { doc.Language = langs[0] } } // Add a file which only occurs in certain branches. func (b *ShardBuilder) Add(doc Document) error { // Skip binary check if already computed (e.g., by Builder.Add // which calls DocChecker.Check before docs reach buildShard). if doc.Category == FileCategoryMissing { if index := bytes.IndexByte(doc.Content, 0); index > 0 { doc.SkipReason = SkipReasonBinary } } if doc.SkipReason != SkipReasonNone { doc.Content = []byte(notIndexedMarker + doc.SkipReason.explanation()) doc.Symbols = nil doc.SymbolsMetaData = nil } DetermineLanguageIfUnknown(&doc) if doc.Category == FileCategoryMissing { DetermineFileCategory(&doc) } sort.Sort(symbolSlice{doc.Symbols, doc.SymbolsMetaData}) var last DocumentSection for i, s := range doc.Symbols { if i > 0 { if last.End > s.Start { return fmt.Errorf("sections overlap") } } last = s } if last.End > uint32(len(doc.Content)) { return fmt.Errorf("section goes past end of content") } if doc.SubRepositoryPath != "" { rel, err := filepath.Rel(doc.SubRepositoryPath, doc.Name) if err != nil || rel == doc.Name { return fmt.Errorf("path %q must start subrepo path %q", doc.Name, doc.SubRepositoryPath) } } docStr, runeSecs, err := b.contentPostings.newSearchableString(doc.Content, doc.Symbols) if err != nil { return err } nameStr, _, err := b.namePostings.newSearchableString([]byte(doc.Name), nil) if err != nil { return err } b.addSymbols(doc.SymbolsMetaData) repoIdx := len(b.repoList) - 1 subRepoIdx, ok := b.subRepoIndices[repoIdx][doc.SubRepositoryPath] if !ok { return fmt.Errorf("unknown subrepo path %q", doc.SubRepositoryPath) } var mask uint64 for _, br := range doc.Branches { m := b.branchMask(br) if m == 0 { return fmt.Errorf("no branch found for %s", br) } mask |= m } if repoIdx > 1<<16 { return fmt.Errorf("too many repos in shard: max is %d", 1<<16) } b.subRepos = append(b.subRepos, subRepoIdx) b.repos = append(b.repos, uint16(repoIdx)) hasher := crc64.New(crc64.MakeTable(crc64.ISO)) hasher.Write(doc.Content) b.contentStrings = append(b.contentStrings, docStr) b.runeDocSections = append(b.runeDocSections, runeSecs...) b.nameStrings = append(b.nameStrings, nameStr) b.docSections = append(b.docSections, doc.Symbols) b.fileEndSymbol = append(b.fileEndSymbol, uint32(len(b.runeDocSections))) b.branchMasks = append(b.branchMasks, mask) b.checksums = append(b.checksums, hasher.Sum(nil)...) langCode, ok := b.languageMap[doc.Language] if !ok { if len(b.languageMap) >= 65535 { return fmt.Errorf("too many languages") } langCode = uint16(len(b.languageMap)) b.languageMap[doc.Language] = langCode } b.languages = append(b.languages, uint8(langCode), uint8(langCode>>8)) category, err := doc.Category.encode() if err != nil { return err } b.categories = append(b.categories, category) return nil } func (b *ShardBuilder) branchMask(br string) uint64 { for i, b := range b.repoList[len(b.repoList)-1].Branches { if b.Name == br { return uint64(1) << uint(i) } } return 0 } // repoIDs returns a list of sourcegraph IDs for the indexed repos. If the ID // is missing or there are no repos, this returns false. func (b *ShardBuilder) repoIDs() ([]uint32, bool) { if len(b.repoList) == 0 { return nil, false } ids := make([]uint32, 0, len(b.repoList)) for _, repo := range b.repoList { if repo.ID == 0 { return nil, false } ids = append(ids, repo.ID) } return ids, true } type DocChecker struct { // A map to count the unique trigrams in a doc. Reused across docs to cut down on allocations. trigrams map[ngram]struct{} } // Check returns a reason why the given contents are probably not source texts. func (t *DocChecker) Check(content []byte, maxTrigramCount int, allowLargeFile bool) SkipReason { if len(content) == 0 { return SkipReasonNone } if len(content) < ngramSize { return SkipReasonTooSmall } if index := bytes.IndexByte(content, 0); index > 0 { return SkipReasonBinary } // PERF: we only need to do the trigram check if the upperbound on content is greater than // our threshold. Also skip the trigram check if the file is explicitly marked as allowed. if trigramsUpperBound := len(content) - ngramSize + 1; trigramsUpperBound <= maxTrigramCount || allowLargeFile { return SkipReasonNone } var cur [3]rune byteCount := 0 t.clearTrigrams(maxTrigramCount) for len(content) > 0 { r, sz := utf8.DecodeRune(content) content = content[sz:] byteCount += sz cur[0], cur[1], cur[2] = cur[1], cur[2], r if cur[0] == 0 { // start of file. continue } t.trigrams[runesToNGram(cur)] = struct{}{} if len(t.trigrams) > maxTrigramCount { // probably not text. return SkipReasonTooManyTrigrams } } return SkipReasonNone } func (t *DocChecker) clearTrigrams(maxTrigramCount int) { if t.trigrams == nil { t.trigrams = make(map[ngram]struct{}, maxTrigramCount) } for key := range t.trigrams { delete(t.trigrams, key) } } // shardName returns the name of the shard for the given prefix, version, and // shard number. func shardName(indexDir string, prefix string, version, n int) string { prefix = url.QueryEscape(prefix) if len(prefix) > 200 { prefix = prefix[:200] + hashString(prefix)[:8] } return filepath.Join(indexDir, fmt.Sprintf("%s_v%d.%05d.zoekt", prefix, version, n)) }