pkg/codesearch: support finding field reads/writes

Author: dvyukov

pkg/aflow/tool/codesearcher/codesearcher.go

@@ -48,7 +48,10 @@ 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.
+If can be used to find all calls or other uses of the given function,
+definition of the given struct/union/enum,
+or all reads/writes of the given struct/union field.
+To find field references use 'struct_name::field_name' syntax.
 `),
 }
 

pkg/codesearch/codesearch.go

@@ -259,8 +259,14 @@ type ReferenceInfo struct {
 
 func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextLines, outputLimit int) (
 	[]ReferenceInfo, int, error) {
+	// Just in case LLM decides to reference structs/fields with the tag.
+	name = strings.TrimSpace(strings.TrimPrefix(strings.TrimPrefix(strings.TrimSpace(name),
+		"struct "), "union "))
+	// We don't export each field as a separate definition,
+	// so we do just name-based match for them.
+	isField := strings.Contains(name, "::")
 	target := index.findDefinition(contextFile, name)
-	if target == nil {
+	if target == nil && !isField {
 		return nil, 0, aflow.BadCallError("requested entity does not exist")
 	}
 	if srcPrefix != "" {
@@ -278,8 +284,8 @@ func (index *Index) FindReferences(contextFile, name, srcPrefix string, contextL
 			// 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 {
+			if ref.Name != name || !isField && (ref.EntityKind != target.Kind ||
+				target.IsStatic && target.Body.File != def.Body.File) {
 				continue
 			}
 			totalCount++

pkg/codesearch/codesearch_test.go

@@ -54,8 +54,8 @@ func testCommand(t *testing.T, index *Index, covered map[string]bool, file strin
 	cmd := fields[0]
 	var args []string
 	for _, arg := range fields[1:] {
-		if arg == `""` {
-			arg = ""
+		if len(arg) >= 2 && arg[0] == '"' && arg[len(arg)-1] == '"' {
+			arg = arg[1 : len(arg)-1]
 		}
 		args = append(args, arg)
 	}

pkg/codesearch/database.go

@@ -57,6 +57,7 @@ const (
 	EntityKindMacro
 	EntityKindEnum
 	EntityKindTypedef
+	EntityKindField
 	entityKindLast
 )
 
@@ -68,6 +69,7 @@ var entityKindNames = [...]string{
 	EntityKindMacro:    "macro",
 	EntityKindEnum:     "enum",
 	EntityKindTypedef:  "typedef",
+	EntityKindField:    "field",
 }
 
 var entityKindBytes = func() [entityKindLast][]byte {
@@ -103,13 +105,17 @@ const (
 	refKindInvalid RefKind = iota
 	RefKindUses
 	RefKindCall
+	RefKindRead
+	RefKindWrite
 	RefKindTakesAddr
 	refKindLast
 )
 
 var refKindNames = [...]string{
 	RefKindUses:      "uses",
 	RefKindCall:      "calls",
+	RefKindRead:      "reads",
+	RefKindWrite:     "writes",
 	RefKindTakesAddr: "takes-address-of",
 }
 

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

@@ -5,4 +5,5 @@ union some_union is defined in source0.h:
   40:	union some_union {
   41:		int x;
   42:		void* p;
-  43:	};
+  43:		struct some_struct s;
+  44:	};

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

@@ -2,7 +2,7 @@ def-source source0.c some_enum yes
 
 enum some_enum is defined in source0.h:
 
-  45:	enum some_enum {
-  46:		enum_foo = 1,
-  47:		enum_bar = 2,
-  48:	};
+  46:	enum some_enum {
+  47:		enum_foo = 1,
+  48:		enum_bar = 2,
+  49:	};

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

@@ -2,4 +2,4 @@ def-source source0.c some_enum_t yes
 
 typedef some_enum_t is defined in source0.h:
 
-  50:	typedef enum some_enum some_enum_t;
+  51:	typedef enum some_enum some_enum_t;

pkg/codesearch/testdata/query-file-index-source

@@ -3,6 +3,7 @@ file-index source0.c
 file source0.c defines the following entities:
 
 function close
+function field_refs
 function func_accepting_a_struct
 function function_with_comment_in_header
 function function_with_quotes_in_type

pkg/codesearch/testdata/query-find-references-field0

@@ -0,0 +1,34 @@
+find-references source0.c some_struct::x "" 1 10
+
+some_struct::x has 5 references:
+
+function field_refs writes it at source0.c:41
+  40:	{
+  41:		p->x = p->y;
+  42:		*(&p->x) = 1;
+
+
+function field_refs takes-address-of it at source0.c:42
+  41:		p->x = p->y;
+  42:		*(&p->x) = 1;
+  43:		u->p = 0;
+
+
+function field_refs writes it at source0.c:44
+  43:		u->p = 0;
+  44:		u->s.x = 2;
+  45:		return p->x;
+
+
+function field_refs reads it at source0.c:45
+  44:		u->s.x = 2;
+  45:		return p->x;
+  46:	}
+
+
+function func_accepting_a_struct reads it at source0.c:31
+  30:	{
+  31:		return ((some_struct_t*)p)->x +
+  32:		       ((union some_union*)p)->x;
+
+

pkg/codesearch/testdata/query-find-references-field1

@@ -0,0 +1,10 @@
+find-references source0.c some_struct::y "" 1 10
+
+some_struct::y has 1 references:
+
+function field_refs reads it at source0.c:41
+  40:	{
+  41:		p->x = p->y;
+  42:		*(&p->x) = 1;
+
+