pkg/codesearch: support searching for references

Extend codesearch clang tool to export info about function references
(calls, takes-address-of).
Add pkg/codesearch command find-references.
Export find-references in pkg/aflow/tools/codesearcher to LLMs.

Update #6469

Author: dvyukov

pkg/aflow/tool/codesearcher/codesearcher.go

@@ -44,6 +44,11 @@ Tool provides full source code for an entity with the given name.
 Entity can be function, struct, or global variable.
 Use it to understand implementation details of an entity.
 For example, how a function works, what precondition error checks it has, etc.
+`),
+	aflow.NewFuncTool("codesearch-find-references", findReferences, `
+Tool finds and lists all references to (uses of) the given entity.
+Entity can be function, struct, or global variable.
+If can be used to find all calls or other uses of the given function.
 `),
 }
 
@@ -95,8 +100,8 @@ type indexEntity struct {
 
 // nolint: lll
 type defCommentArgs struct {
-	SourceFile string `jsonschema:"Source file path that references the entity. It helps to restrict scope of the search, if there are different definitions with the same name in different source files."`
-	Name       string `jsonschema:"Name of the entity of interest."`
+	ContextFile string `jsonschema:"Source file path that references the entity. It helps to restrict scope of the search, if there are different definitions with the same name in different source files."`
+	Name        string `jsonschema:"Name of the entity of interest."`
 }
 
 type defCommentResult struct {
@@ -106,7 +111,7 @@ type defCommentResult struct {
 
 // nolint: lll
 type defSourceArgs struct {
-	SourceFile   string `jsonschema:"Source file path that references the entity. It helps to restrict scope of the search, if there are different definitions with the same name in different source files."`
+	ContextFile  string `jsonschema:"Source file path that references the entity. It helps to restrict scope of the search, if there are different definitions with the same name in different source files."`
 	Name         string `jsonschema:"Name of the entity of interest."`
 	IncludeLines bool   `jsonschema:"Whether to include line numbers in the output or not. Line numbers may distract you, so ask for them only if you need to match lines elsewhere with the source code."`
 }
@@ -181,7 +186,7 @@ func fileIndex(ctx *aflow.Context, state prepareResult, args fileIndexArgs) (fil
 }
 
 func definitionComment(ctx *aflow.Context, state prepareResult, args defCommentArgs) (defCommentResult, error) {
-	info, err := state.Index.DefinitionComment(args.SourceFile, args.Name)
+	info, err := state.Index.DefinitionComment(args.ContextFile, args.Name)
 	if err != nil {
 		return defCommentResult{}, err
 	}
@@ -192,7 +197,7 @@ func definitionComment(ctx *aflow.Context, state prepareResult, args defCommentA
 }
 
 func definitionSource(ctx *aflow.Context, state prepareResult, args defSourceArgs) (defSourceResult, error) {
-	info, err := state.Index.DefinitionSource(args.SourceFile, args.Name, args.IncludeLines)
+	info, err := state.Index.DefinitionSource(args.ContextFile, args.Name, args.IncludeLines)
 	if err != nil {
 		return defSourceResult{}, err
 	}
@@ -201,3 +206,56 @@ func definitionSource(ctx *aflow.Context, state prepareResult, args defSourceArg
 		SourceCode: info.Body,
 	}, nil
 }
+
+// nolint: lll
+type findReferencesArgs struct {
+	ContextFile         string `jsonschema:"Source file path that references the entity. It helps to restrict scope of the search, if there are different definitions with the same name in different source files." json:",omitempty"`
+	Name                string `jsonschema:"Name of the entity of interest."`
+	SourceTreePrefix    string `jsonschema:"Prefix of the sources tree where to search for references. Can be used to restrict search to e.g. net/ipv4/. Pass an empty string to find all references." json:",omitempty"`
+	IncludeSnippetLines uint   `jsonschema:"If set to non-0, output will include source code snippets with that many lines of context. If set to 0, no source snippets will be included. Snippets only show the referencing entity, so to see e.g. whole referencing functions pass a large value, e.g. 10000" json:",omitempty"`
+}
+
+// nolint: lll
+type findReferencesResult struct {
+	TruncatedOutput bool        `jsonschema:"Set if there were too many references, and the output is truncated. If you get truncated output, you may try to either request w/o source code snippets by passing IncludeSnippetLines=0 (which has higher limit on the number of output references), or restrict search to some prefix of the source tree with SourceTreePrefix argument."`
+	References      []Reference `jsonschema:"List of requested references."`
+}
+
+type Reference struct {
+	ReferencingEntityKind string `jsonschema:"Kind of the referencing entity (function, struct, etc)."`
+	ReferencingEntityName string `jsonschema:"Name of the referencing entity."`
+	ReferenceKind         string `jsonschema:"Kind of the reference (calls, takes-address, reads, writes-to, etc)."`
+	SourceFile            string `jsonschema:"Source file of the reference."`
+	SourceLine            int    `jsonschema:"Source line of the reference."`
+	SourceSnippet         string `jsonschema:"Surrounding code snippet, if requested." json:",omitempty"`
+}
+
+func findReferences(ctx *aflow.Context, state prepareResult, args findReferencesArgs) (findReferencesResult, error) {
+	outputLimit := 20
+	if args.IncludeSnippetLines == 0 {
+		outputLimit = 1000
+	} else if args.IncludeSnippetLines < 10 {
+		outputLimit = 100
+	}
+	refs, totalCount, err := state.Index.FindReferences(
+		args.ContextFile, args.Name, args.SourceTreePrefix,
+		int(args.IncludeSnippetLines), outputLimit)
+	if err != nil {
+		return findReferencesResult{}, err
+	}
+	var results []Reference
+	for _, ref := range refs {
+		results = append(results, Reference{
+			ReferencingEntityKind: ref.ReferencingEntityKind,
+			ReferencingEntityName: ref.ReferencingEntityName,
+			ReferenceKind:         ref.ReferenceKind,
+			SourceFile:            ref.SourceFile,
+			SourceLine:            ref.SourceLine,
+			SourceSnippet:         ref.SourceSnippet,
+		})
+	}
+	return findReferencesResult{
+		TruncatedOutput: totalCount > len(refs),
+		References:      results,
+	}, nil
+}

pkg/clangtool/tooltest/tooltest.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"testing"
@@ -82,7 +83,13 @@ func ForEachTestFile(t *testing.T, fn func(t *testing.T, cfg *clangtool.Config,
 }
 
 func forEachTestFile(t *testing.T, fn func(t *testing.T, file string)) {
-	files, err := filepath.Glob(filepath.Join(osutil.Abs("testdata"), "*.c"))
+	var files []string
+	err := filepath.WalkDir(osutil.Abs("testdata"), func(path string, d fs.DirEntry, err error) error {
+		if d.Name()[0] != '.' && filepath.Ext(d.Name()) == ".c" {
+			files = append(files, path)
+		}
+		return err
+	})
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/codesearch/codesearch.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	"strconv"
 	"strings"
 	"syscall"
 
@@ -80,6 +81,28 @@ var Commands = []Command{
 		}
 		return fmt.Sprintf("%v %v is defined in %v:\n\n%v", info.Kind, args[1], info.File, info.Body), nil
 	}},
+	{"find-references", 5, func(index *Index, args []string) (string, error) {
+		contextLines, err := strconv.Atoi(args[3])
+		if err != nil {
+			return "", fmt.Errorf("failed to parse number of context lines %q: %v", args[3], err)
+		}
+		outputLimit, err := strconv.Atoi(args[4])
+		if err != nil {
+			return "", fmt.Errorf("failed to parse output limit %q: %v", args[4], err)
+		}
+		refs, totalCount, err := index.FindReferences(args[0], args[1], args[2], contextLines, outputLimit)
+		if err != nil {
+			return "", err
+		}
+		b := new(strings.Builder)
+		fmt.Fprintf(b, "%v has %v references:\n\n", args[1], totalCount)
+		for _, ref := range refs {
+			fmt.Fprintf(b, "%v %v %v it at %v:%v\n%v\n\n",
+				ref.ReferencingEntityKind, ref.ReferencingEntityName, ref.ReferenceKind,
+				ref.SourceFile, ref.SourceLine, ref.SourceSnippet)
+		}
+		return b.String(), nil
+	}},
 }
 
 func IsSourceFile(file string) bool {
@@ -225,6 +248,68 @@ func (index *Index) definitionSource(contextFile, name string, comment, includeL
 	}, nil
 }
 
+type ReferenceInfo struct {
+	ReferencingEntityKind string
+	ReferencingEntityName string
+	ReferenceKind         string
+	SourceFile            string
+	SourceLine            int
+	SourceSnippet         string
+}
+
+func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextLines, outputLimit int) ([]*ReferenceInfo, int, error) {
+	target := index.findDefinition(contextFile, name)
+	if target == nil {
+		return nil, 0, aflow.BadCallError("requested entity does not exist")
+	}
+	if srcPrefix != "" {
+		srcPrefix = filepath.Clean(srcPrefix)
+	}
+	totalCount := 0
+	var results []*ReferenceInfo
+	for _, def := range index.db.Definitions {
+		if !strings.HasPrefix(def.Body.File, srcPrefix) {
+			continue
+		}
+		for _, ref := range def.Refs {
+			// TODO: this mis-handles the following case:
+			// the target is a non-static 'foo' in some file,
+			// the reference is in another file and refers to a static 'foo'
+			// defined in that file (which is not the target 'foo').
+			if !(ref.EntityKind == target.Kind && ref.Name == target.Name &&
+				(!target.IsStatic || target.Body.File == def.Body.File)) {
+				continue
+			}
+			totalCount++
+			if totalCount > outputLimit {
+				continue
+			}
+			snippet := ""
+			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),
+				}
+				var err error
+				snippet, err = index.formatSource(lines, true)
+				if err != nil {
+					return nil, 0, err
+				}
+			}
+			results = append(results, &ReferenceInfo{
+				ReferencingEntityKind: def.Kind,
+				ReferencingEntityName: def.Name,
+				ReferenceKind:         ref.Kind,
+				SourceFile:            def.Body.File,
+				SourceLine:            ref.Line,
+				SourceSnippet:         snippet,
+			})
+		}
+	}
+	return results, totalCount, nil
+}
+
 func (index *Index) findDefinition(contextFile, name string) *Definition {
 	var weakMatch *Definition
 	for _, def := range index.db.Definitions {
@@ -269,7 +354,7 @@ func formatSourceFile(file string, start, end int, includeLines bool) (string, e
 	b := new(strings.Builder)
 	for line := start; line <= end; line++ {
 		if includeLines {
-			fmt.Fprintf(b, "%4v:\t%s\n", line, lines[line])
+			fmt.Fprintf(b, "%4v:\t%s\n", line+1, lines[line])
 		} else {
 			fmt.Fprintf(b, "%s\n", lines[line])
 		}

pkg/codesearch/codesearch_test.go

@@ -47,16 +47,24 @@ func testCommand(t *testing.T, index *Index, covered map[string]bool, file strin
 		t.Fatal(err)
 	}
 	query, _, _ := bytes.Cut(data, []byte{'\n'})
-	args := strings.Fields(string(query))
-	if len(args) == 0 {
+	fields := strings.Fields(string(query))
+	if len(fields) == 0 {
 		t.Fatal("no command found")
 	}
-	result, err := index.Command(args[0], args[1:])
+	cmd := fields[0]
+	var args []string
+	for _, arg := range fields[1:] {
+		if arg == `""` {
+			arg = ""
+		}
+		args = append(args, arg)
+	}
+	result, err := index.Command(cmd, args)
 	if err != nil {
 		// This is supposed to test aflow.BadCallError messages.
 		result = err.Error() + "\n"
 	}
-	got := append([]byte(strings.Join(args, " ")+"\n\n"), result...)
+	got := append([]byte(strings.Join(fields, " ")+"\n\n"), result...)
 	tooltest.CompareGoldenData(t, file, got)
-	covered[args[0]] = true
+	covered[cmd] = true
 }

pkg/codesearch/database.go

@@ -14,12 +14,20 @@ type Database struct {
 }
 
 type Definition struct {
-	Kind     string    `json:"kind,omitempty"`
-	Name     string    `json:"name,omitempty"`
-	Type     string    `json:"type,omitempty"`
-	IsStatic bool      `json:"is_static,omitempty"`
-	Body     LineRange `json:"body,omitempty"`
-	Comment  LineRange `json:"comment,omitempty"`
+	Kind     string      `json:"kind,omitempty"`
+	Name     string      `json:"name,omitempty"`
+	Type     string      `json:"type,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"`
 }
 
 type LineRange struct {

pkg/codesearch/testdata/mm/refs.c

@@ -0,0 +1,6 @@
+int refs2();
+
+void ref_in_mm()
+{
+	refs2();
+}

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

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

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

@@ -0,0 +1 @@
+{}

pkg/codesearch/testdata/query-def-source-header

@@ -2,7 +2,7 @@ def-source source0.c function_with_comment_in_header yes
 
 function function_with_comment_in_header is defined in source0.c:
 
-  18:	void function_with_comment_in_header()
-  19:	{
-  20:		same_name_in_several_files();
-  21:	}
+  19:	void function_with_comment_in_header()
+  20:	{
+  21:		same_name_in_several_files();
+  22:	}

pkg/codesearch/testdata/query-def-source-open

@@ -2,10 +2,10 @@ def-source source0.c open yes
 
 function open is defined in source0.c:
 
-   5:	/*
-   6:	 * Comment about open.
-   7:	 */
-   8:	int open()
-   9:	{
-  10:		return 0;
-  11:	}
+   6:	/*
+   7:	 * Comment about open.
+   8:	 */
+   9:	int open()
+  10:	{
+  11:		return 0;
+  12:	}