dvyukov some ai changes

pkg/clangtool/clangtool.go

@@ -32,7 +32,7 @@ type Config struct {
 
 type OutputDataPtr[T any] interface {
 	*T
-	Merge(*T)
+	Merge(*T, *Verifier)
 	SetSourceFile(string, func(filename string) string)
 	Finalize(*Verifier)
 }
@@ -73,21 +73,22 @@ func Run[Output any, OutputPtr OutputDataPtr[Output]](cfg *Config) (OutputPtr, e
 	}
 	close(files)
 
+	v := NewVerifier(cfg.KernelSrc, cfg.KernelObj)
 	out := OutputPtr(new(Output))
 	for range cmds {
 		res := <-results
 		if res.err != nil {
 			return nil, res.err
 		}
-		out.Merge(res.out)
+		out.Merge(res.out, v)
 	}
 	// Finalize the output (sort, dedup, etc), and let the output verify
 	// that all source file names, line numbers, etc are valid/present.
 	// If there are any bogus entries, it's better to detect them early,
 	// than to crash/error much later when the info is used.
 	// Some of the source files (generated) may be in the obj dir.
-	srcDirs := []string{cfg.KernelSrc, cfg.KernelObj}
-	if err := Finalize(out, srcDirs); err != nil {
+	out.Finalize(v)
+	if err := v.Error(); err != nil {
 		return nil, err
 	}
 	if cfg.CacheFile != "" {
@@ -103,24 +104,26 @@ func Run[Output any, OutputPtr OutputDataPtr[Output]](cfg *Config) (OutputPtr, e
 	return out, nil
 }
 
-func Finalize[Output any, OutputPtr OutputDataPtr[Output]](out OutputPtr, srcDirs []string) error {
-	v := &Verifier{
-		srcDirs:   srcDirs,
+type Verifier struct {
+	srcDirs   []string
+	fileCache map[string]int // file->line count (-1 is cached for missing files)
+	err       strings.Builder
+}
+
+func NewVerifier(src ...string) *Verifier {
+	return &Verifier{
+		srcDirs:   src,
 		fileCache: make(map[string]int),
 	}
-	out.Finalize(v)
+}
+
+func (v *Verifier) Error() error {
 	if v.err.Len() == 0 {
 		return nil
 	}
 	return errors.New(v.err.String())
 }
 
-type Verifier struct {
-	srcDirs   []string
-	fileCache map[string]int // file->line count (-1 is cached for missing files)
-	err       strings.Builder
-}
-
 func (v *Verifier) Filename(file string) {
 	if _, ok := v.fileCache[file]; ok {
 		return

pkg/clangtool/tooltest/tooltest.go

@@ -42,14 +42,16 @@ func TestClangTool[Output any, OutputPtr clangtool.OutputDataPtr[Output]](t *tes
 
 func LoadOutput[Output any, OutputPtr clangtool.OutputDataPtr[Output]](t *testing.T) OutputPtr {
 	out := OutputPtr(new(Output))
+	v := clangtool.NewVerifier("testdata")
 	forEachTestFile(t, func(t *testing.T, file string) {
 		tmp, err := osutil.ReadJSON[OutputPtr](file + ".json")
 		if err != nil {
 			t.Fatal(err)
 		}
-		out.Merge(tmp)
+		out.Merge(tmp, v)
 	})
-	if err := clangtool.Finalize(out, []string{"testdata"}); err != nil {
+	out.Finalize(v)
+	if err := v.Error(); err != nil {
 		t.Fatal(err)
 	}
 	return out

pkg/codesearch/codesearch.go

@@ -206,7 +206,7 @@ func (index *Index) FileIndex(file string) ([]Entity, error) {
 	for _, def := range index.db.Definitions {
 		if def.Body.File == file {
 			entities = append(entities, Entity{
-				Kind: def.Kind,
+				Kind: def.Kind.String(),
 				Name: def.Name,
 			})
 		}
@@ -243,7 +243,7 @@ func (index *Index) definitionSource(contextFile, name string, comment, includeL
 	}
 	return &EntityInfo{
 		File: def.Body.File,
-		Kind: def.Kind,
+		Kind: def.Kind.String(),
 		Body: src,
 	}, nil
 }
@@ -266,6 +266,7 @@ func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextL
 	if srcPrefix != "" {
 		srcPrefix = filepath.Clean(srcPrefix)
 	}
+	contextLines = min(contextLines, 10000)
 	totalCount := 0
 	var results []ReferenceInfo
 	for _, def := range index.db.Definitions {
@@ -289,8 +290,8 @@ func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextL
 			if contextLines > 0 {
 				lines := LineRange{
 					File:      def.Body.File,
-					StartLine: max(def.Body.StartLine, ref.Line-contextLines),
-					EndLine:   min(def.Body.EndLine, ref.Line+contextLines),
+					StartLine: max(def.Body.StartLine, uint32(max(0, int(ref.Line)-contextLines))),
+					EndLine:   min(def.Body.EndLine, ref.Line+uint32(contextLines)),
 				}
 				var err error
 				snippet, err = index.formatSource(lines, true)
@@ -299,11 +300,11 @@ func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextL
 				}
 			}
 			results = append(results, ReferenceInfo{
-				ReferencingEntityKind: def.Kind,
+				ReferencingEntityKind: def.Kind.String(),
 				ReferencingEntityName: def.Name,
-				ReferenceKind:         ref.Kind,
+				ReferenceKind:         ref.Kind.String(),
 				SourceFile:            def.Body.File,
-				SourceLine:            ref.Line,
+				SourceLine:            int(ref.Line),
 				SourceSnippet:         snippet,
 			})
 		}
@@ -342,7 +343,7 @@ func (index *Index) formatSource(lines LineRange, includeLines bool) (string, er
 		if !osutil.IsExist(file) {
 			continue
 		}
-		return formatSourceFile(file, lines.StartLine, lines.EndLine, includeLines)
+		return formatSourceFile(file, int(lines.StartLine), int(lines.EndLine), includeLines)
 	}
 	return "", fmt.Errorf("codesearch: can't find %q file in any of %v", lines.File, index.srcDirs)
 }

pkg/codesearch/database.go

@@ -4,6 +4,10 @@
 package codesearch
 
 import (
+	"bytes"
+	"fmt"
+	"maps"
+	"slices"
 	"strings"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -13,29 +17,127 @@ import (
 
 type Database struct {
 	Definitions []*Definition `json:"definitions,omitempty"`
+
+	mergeCache   map[string]*Definition
+	reverseCache map[*Definition]string
+	stringCache  map[string]string
 }
 
 type Definition struct {
-	Kind     string      `json:"kind,omitempty"`
 	Name     string      `json:"name,omitempty"`
 	Type     string      `json:"type,omitempty"`
+	Kind     EntityKind  `json:"kind,omitempty"`
 	IsStatic bool        `json:"is_static,omitempty"`
 	Body     LineRange   `json:"body,omitempty"`
 	Comment  LineRange   `json:"comment,omitempty"`
 	Refs     []Reference `json:"refs,omitempty"`
 }
 
 type Reference struct {
-	Kind       string `json:"kind,omitempty"`
-	EntityKind string `json:"entity_kind,omitempty"`
-	Name       string `json:"name,omitempty"`
-	Line       int    `json:"line,omitempty"`
+	Name       string     `json:"name,omitempty"`
+	Kind       RefKind    `json:"kind,omitempty"`
+	EntityKind EntityKind `json:"entity_kind,omitempty"`
+	Line       uint32     `json:"line,omitempty"`
 }
 
 type LineRange struct {
 	File      string `json:"file,omitempty"`
-	StartLine int    `json:"start_line,omitempty"`
-	EndLine   int    `json:"end_line,omitempty"`
+	StartLine uint32 `json:"start_line,omitempty"`
+	EndLine   uint32 `json:"end_line,omitempty"`
+}
+
+type EntityKind uint8
+
+const (
+	entityKindInvalid EntityKind = iota
+	EntityKindFunction
+	EntityKindStruct
+	EntityKindUnion
+	EntityKindVariable
+	EntityKindMacro
+	EntityKindEnum
+	EntityKindTypedef
+	entityKindLast
+)
+
+var entityKindNames = [...]string{
+	EntityKindFunction: "function",
+	EntityKindStruct:   "struct",
+	EntityKindUnion:    "union",
+	EntityKindVariable: "variable",
+	EntityKindMacro:    "macro",
+	EntityKindEnum:     "enum",
+	EntityKindTypedef:  "typedef",
+}
+
+var entityKindBytes = func() [entityKindLast][]byte {
+	var ret [entityKindLast][]byte
+	for k, v := range entityKindNames {
+		ret[k] = []byte("\"" + v + "\"")
+	}
+	return ret
+}()
+
+func (v *EntityKind) String() string {
+	return entityKindNames[*v]
+}
+
+func (v *EntityKind) MarshalJSON() ([]byte, error) {
+	return entityKindBytes[*v], nil
+}
+
+func (v *EntityKind) UnmarshalJSON(data []byte) error {
+	*v = entityKindInvalid
+	for k, val := range entityKindBytes {
+		if bytes.Equal(data, val) {
+			*v = EntityKind(k)
+			break
+		}
+	}
+	return nil
+}
+
+type RefKind uint8
+
+const (
+	refKindInvalid RefKind = iota
+	RefKindUses
+	RefKindCall
+	RefKindTakesAddr
+	refKindLast
+)
+
+var refKindNames = [...]string{
+	RefKindUses:      "uses",
+	RefKindCall:      "calls",
+	RefKindTakesAddr: "takes-address-of",
+}
+
+var refKindBytes = func() [refKindLast][]byte {
+	var ret [refKindLast][]byte
+	for k, v := range refKindNames {
+		ret[k] = []byte("\"" + v + "\"")
+	}
+	return ret
+}()
+
+func (v *RefKind) String() string {
+	return refKindNames[*v]
+}
+
+func (v *RefKind) MarshalJSON() ([]byte, error) {
+	return refKindBytes[*v], nil
+}
+
+func (v *RefKind) UnmarshalJSON(data []byte) error {
+	*v = refKindInvalid
+	for k, val := range refKindBytes {
+		if bytes.Equal(data, val) {
+			*v = RefKind(k)
+			break
+		}
+	}
+	return nil
 }
 
 // DatabaseFormatHash contains a hash uniquely identifying format of the database.
@@ -44,29 +146,50 @@ type LineRange struct {
 var DatabaseFormatHash = func() string {
 	// Semantic version should be bumped when the schema does not change,
 	// but stored values changes.
-	const semanticVersion = "2"
+	const semanticVersion = "3"
 	schema, err := jsonschema.For[Database](nil)
 	if err != nil {
 		panic(err)
 	}
 	return hash.String(schema, semanticVersion)
 }()
 
-func (db *Database) Merge(other *Database) {
-	db.Definitions = append(db.Definitions, other.Definitions...)
-}
-
-func (db *Database) Finalize(v *clangtool.Verifier) {
-	db.Definitions = clangtool.SortAndDedupSlice(db.Definitions)
-
-	for _, def := range db.Definitions {
-		v.LineRange(def.Body.File, def.Body.StartLine, def.Body.EndLine)
+func (db *Database) Merge(other *Database, v *clangtool.Verifier) {
+	if db.mergeCache == nil {
+		db.mergeCache = make(map[string]*Definition)
+		db.reverseCache = make(map[*Definition]string)
+		db.stringCache = make(map[string]string)
+	}
+	for _, def := range other.Definitions {
+		id := fmt.Sprintf("%v-%v-%v", def.Kind, def.Name, def.Body.File)
+		if _, ok := db.mergeCache[id]; ok {
+			continue
+		}
+		db.mergeCache[id] = def
+		db.reverseCache[def] = id
+		v.LineRange(def.Body.File, int(def.Body.StartLine), int(def.Body.EndLine))
 		if def.Comment.File != "" {
-			v.LineRange(def.Comment.File, def.Comment.StartLine, def.Comment.EndLine)
+			v.LineRange(def.Comment.File, int(def.Comment.StartLine), int(def.Comment.EndLine))
+		}
+		db.intern(&def.Name)
+		db.intern(&def.Type)
+		db.intern(&def.Body.File)
+		db.intern(&def.Comment.File)
+		for _, ref := range def.Refs {
+			db.intern(&ref.Name)
 		}
 	}
 }
 
+func (db *Database) Finalize(v *clangtool.Verifier) {
+	db.Definitions = slices.Collect(maps.Values(db.mergeCache))
+	slices.SortFunc(db.Definitions, func(a, b *Definition) int {
+		return strings.Compare(db.reverseCache[a], db.reverseCache[b])
+	})
+	db.mergeCache = nil
+	db.reverseCache = nil
+}
+
 // SetSoureFile attaches the source file to the entities that need it.
 // The clang tool could do it, but it looks easier to do it here.
 func (db *Database) SetSourceFile(file string, updatePath func(string) string) {
@@ -78,3 +201,15 @@ func (db *Database) SetSourceFile(file string, updatePath func(string) string) {
 		}
 	}
 }
+
+func (db *Database) intern(str *string) {
+	if *str == "" {
+		return
+	}
+	v, ok := db.stringCache[*str]
+	if !ok {
+		v = strings.Clone(*str)
+		db.stringCache[v] = v
+	}
+	*str = v
+}

pkg/codesearch/testdata/mm/refs.c.json

@@ -1,9 +1,9 @@
 {
 	"definitions": [
 		{
-			"kind": "function",
 			"name": "ref_in_mm",
 			"type": "void ()",
+			"kind": "function",
 			"body": {
 				"file": "mm/refs.c",
 				"start_line": 3,
@@ -12,9 +12,9 @@
 			"comment": {},
 			"refs": [
 				{
+					"name": "refs2",
 					"kind": "calls",
 					"entity_kind": "function",
-					"name": "refs2",
 					"line": 5
 				}
 			]