rust: index the items written inside item-level macro bodies

strategies/extractors.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	sitter "github.com/smacker/go-tree-sitter"
+	"github.com/smacker/go-tree-sitter/rust"
 
 	"github.com/provasign/astkit"
 	"github.com/provasign/astkit/internalast"
@@ -2420,6 +2421,8 @@ func rustVisit(node *sitter.Node, filePath, blobSHA string, src []byte, imports
 		case "enum_item":
 			rustNamedItem(n, astkit.KindEnum, filePath, blobSHA, src, imports, out)
 			rustEnumVariants(n, src, out)
+		case "macro_invocation":
+			rustMacroItems(n, filePath, blobSHA, src, imports, implType, implBounds, implTrait, out)
 		case "macro_definition":
 			// macro_rules! name { ... }: a named item callers invoke as name!.
 			rustNamedItem(n, astkit.KindMacro, filePath, blobSHA, src, imports, out)
@@ -2504,6 +2507,45 @@ func rustVisit(node *sitter.Node, filePath, blobSHA string, src []byte, imports
 	}
 }
 
+// rustMacroItems indexes the items written out inside an item-level macro
+// invocation: `cfg_rt! { pub fn spawn(...) {...} }`, `pin_project! { pub
+// struct Framed {...} }`. tree-sitter leaves a macro's body as an unparsed
+// token tree, so tokio's cfg-gated API had no symbols at all. The body is
+// re-parsed as Rust over its own byte range (positions stay those of the
+// file) and visited like any item list; nested wrappers recurse on a
+// strictly smaller range. Only brace bodies outside function bodies and
+// macro_rules templates are read: a call-like `m!(a, b)` is arguments, not
+// items, and a name a macro generates (`rgtest!(r16, ..)` -> fn r16) is the
+// compiler's to report.
+func rustMacroItems(n *sitter.Node, filePath, blobSHA string, src []byte, imports []string, implType string, implBounds []string, implTrait string, out *[]astkit.Symbol) {
+	tt := internalast.FindChildByType(n, "token_tree")
+	if tt == nil || tt.EndByte()-tt.StartByte() < 2 || src[tt.StartByte()] != '{' || src[tt.EndByte()-1] != '}' {
+		return
+	}
+	for p := n.Parent(); p != nil; p = p.Parent() {
+		switch p.Type() {
+		case "function_item", "closure_expression", "macro_definition":
+			return
+		}
+	}
+	if len(strings.TrimSpace(string(src[tt.StartByte()+1:tt.EndByte()-1]))) == 0 {
+		return
+	}
+	start, end := tt.StartPoint(), tt.EndPoint()
+	start.Column++
+	end.Column--
+	p := sitter.NewParser()
+	defer p.Close()
+	p.SetLanguage(rust.GetLanguage())
+	p.SetIncludedRanges([]sitter.Range{{StartPoint: start, EndPoint: end, StartByte: tt.StartByte() + 1, EndByte: tt.EndByte() - 1}})
+	tree, err := p.ParseCtx(context.Background(), nil, src)
+	if err != nil || tree == nil {
+		return
+	}
+	defer tree.Close()
+	rustVisit(tree.RootNode(), filePath, blobSHA, src, imports, implType, implBounds, implTrait, out)
+}
+
 // rustScopeQualify prefixes a symbol found inside a module or function body
 // with that scope. A member (field, method, associated item) whose qualified
 // name is still its bare name keeps its owner type in the path: `mod net {

strategies/rust_decl_test.go

@@ -106,3 +106,67 @@ func TestRustUnionIsIndexedWithFields(t *testing.T) {
 		"Bits": astkit.KindStruct, "Bits.i": astkit.KindField, "Bits.f": astkit.KindField,
 	})
 }
+
+func TestRustItemMacroBodiesAreIndexed(t *testing.T) {
+	src := `cfg_rt! {
+    pub fn spawn(task: Task) -> Handle {
+        start(task)
+    }
+    cfg_net! {
+        pub struct TcpStream { fd: i32 }
+    }
+}
+pub struct Runtime;
+impl Runtime {
+    cfg_rt! {
+        pub fn block_on(&self) {}
+    }
+}
+macro_rules! cfg_rt {
+    ($($item:item)*) => { $($item)* };
+}
+rgtest!(r16, |dir: Dir| {
+    dir.create("x");
+});
+fn body() {
+    helper! { struct Hidden; }
+}
+`
+	syms, _ := extract(t, astkit.LangRust, src)
+	gprWantDecls(t, syms, map[string]astkit.SymbolKind{
+		"spawn": astkit.KindFunction, "TcpStream": astkit.KindStruct,
+		"Runtime.block_on": astkit.KindMethod,
+	})
+	for _, s := range syms {
+		switch s.Name {
+		case "spawn":
+			if s.Span.Start != 2 || s.Span.End != 4 || len(s.CallSites) == 0 {
+				t.Errorf("spawn keeps its file position and its calls: %+v", s)
+			}
+		case "fd":
+			if s.ParentName != "TcpStream" || s.Span.Start != 6 {
+				t.Errorf("nested wrapper field: %+v", s)
+			}
+		case "r16", "Hidden", "item":
+			t.Errorf("not an item written in a brace body at item level: %+v", s)
+		}
+	}
+}
+
+// A wrapper body that defines a helper macro_rules (tokio's AsyncReadExt)
+// still declares its items.
+func TestRustItemMacroBodyWithHelperMacro(t *testing.T) {
+	src := `cfg_io_util! {
+    macro_rules! read_impl {
+        ($name:ident) => { fn $name(&mut self) {} };
+    }
+    pub trait AsyncReadExt {
+        fn read(&mut self) {}
+    }
+}
+`
+	syms, _ := extract(t, astkit.LangRust, src)
+	gprWantDecls(t, syms, map[string]astkit.SymbolKind{
+		"AsyncReadExt": astkit.KindTrait, "AsyncReadExt.read": astkit.KindMethod, "read_impl": astkit.KindMacro,
+	})
+}