[WIP] cl: preserve linknames for main test variants

cl/import.go

@@ -343,8 +343,9 @@ func collectDeclarationDirectivesWithOptions(prog llssa.Program, fset *token.Fil
 	return linkCollected, nil
 }
 
-// collectGoLinknames follows cmd/compile's package-scoped go:linkname behavior.
-func collectGoLinknames(prog llssa.Program, comments []*ast.CommentGroup, syms map[string]string) {
+// collectGoLinknames follows cmd/compile's package-scoped go:linkname behavior
+// and adds raw import-path aliases for main-package test variants.
+func collectGoLinknames(prog llssa.Program, fset *token.FileSet, comments []*ast.CommentGroup, syms map[string]string, rawMainPath string) {
 	const prefix = "//go:linkname "
 	for _, group := range comments {
 		for _, comment := range group.List {
@@ -356,12 +357,47 @@ func collectGoLinknames(prog llssa.Program, comments []*ast.CommentGroup, syms m
 				continue
 			}
 			if fullName, ok := syms[fields[0]]; ok {
-				prog.SetLinkname(fullName, strings.Join(fields[1:], " "))
+				link := strings.Join(fields[1:], " ")
+				prog.SetLinkname(fullName, link)
+				if rawMainPath != "" {
+					prog.SetLinkname(rawMainPath+"."+fields[0], rewriteMainLinkname(rawMainPath, link, !isTestSource(fset, comment.Pos())))
+				}
 			}
 		}
 	}
 }
 
+// addMainPackageLinknameAlias records the raw import-path spelling used by
+// test variants. Ordinary main packages use the canonical "main" prefix,
+// while go test emits their symbols under the package import path. The target
+// remains unchanged: test variants do not define main.* symbols.
+func addMainPackageLinknameAlias(prog llssa.Program, rawMainPath, fullName, inPkgName string, rewriteMainTarget bool) {
+	if rawMainPath == "" {
+		return
+	}
+	if link, ok := prog.Linkname(fullName); ok {
+		prog.SetLinkname(rawMainPath+"."+inPkgName, rewriteMainLinkname(rawMainPath, link, rewriteMainTarget))
+	}
+}
+
+func rewriteMainLinkname(rawMainPath, link string, rewrite bool) string {
+	if rewrite && strings.HasPrefix(link, "main.") {
+		return rawMainPath + strings.TrimPrefix(link, "main")
+	}
+	return link
+}
+
+func isTestSource(fset *token.FileSet, pos token.Pos) bool {
+	return fset != nil && strings.HasSuffix(fset.Position(pos).Filename, "_test.go")
+}
+
+func rawMainPathOf(pkg *types.Package) string {
+	if pkg == nil || pkg.Name() != "main" || pkg.Path() == "" || pkg.Path() == llssa.PathOf(pkg) {
+		return ""
+	}
+	return pkg.Path()
+}
+
 func (p *context) processLinknameByDoc(doc *ast.CommentGroup, fullName, inPkgName string, isVar, allowExport bool) bool {
 	if doc != nil {
 		for n := len(doc.List) - 1; n >= 0; n-- {
@@ -839,6 +875,7 @@ func ParsePkgSyntaxWithOptions(prog llssa.Program, fset *token.FileSet, pkg *typ
 	}
 	ctx := &context{prog: prog, options: options}
 	pkgPath := llssa.PathOf(pkg)
+	rawMainPath := rawMainPathOf(pkg)
 	syms := make(map[string]string)
 	var fileComments []*ast.CommentGroup
 	for _, file := range files {
@@ -863,6 +900,7 @@ func ParsePkgSyntaxWithOptions(prog llssa.Program, fset *token.FileSet, pkg *typ
 				if err != nil {
 					return err
 				}
+				addMainPackageLinknameAlias(prog, rawMainPath, fullName, inPkgName, !isTestSource(fset, decl.Pos()))
 				if !hasLinkname && pkg.Name() == "C" && decl.Recv == nil && token.IsExported(inPkgName) {
 					exportName := strings.TrimPrefix(inPkgName, "X")
 					prog.SetLinkname(fullName, exportName)
@@ -882,6 +920,7 @@ func ParsePkgSyntaxWithOptions(prog llssa.Program, fset *token.FileSet, pkg *typ
 							if _, err := collectDeclarationDirectivesWithOptions(prog, fset, decl.Doc, pkgPath+"."+inPkgName, inPkgName, token.NoPos, options); err != nil {
 								return err
 							}
+							addMainPackageLinknameAlias(prog, rawMainPath, pkgPath+"."+inPkgName, inPkgName, !isTestSource(fset, decl.Pos()))
 						}
 					}
 					vars, err := locality.ScanPackageVar(fset, decl)
@@ -902,7 +941,7 @@ func ParsePkgSyntaxWithOptions(prog llssa.Program, fset *token.FileSet, pkg *typ
 			}
 		}
 	}
-	collectGoLinknames(prog, fileComments, syms)
+	collectGoLinknames(prog, fset, fileComments, syms, rawMainPath)
 	prog.MarkPackageSyntaxParsed(pkg)
 	return nil
 }

cl/import_coverage_test.go

@@ -278,6 +278,67 @@ func TestParsePkgSyntaxCollectsLinknames(t *testing.T) {
 	if got, ok := prog.Linkname("example.com/p.alias"); !ok || got != "C.alias" {
 		t.Fatalf("cross-file linkname = (%q,%v), want (C.alias,true)", got, ok)
 	}
+
+	mainSrc := []string{
+		"package main\nimport _ \"unsafe\"\n",
+		"package main\n//go:linkname demo6 main.demo\nfunc demo6()\n",
+	}
+	mainFset := token.NewFileSet()
+	mainFiles := make([]*ast.File, 0, len(mainSrc))
+	for i, src := range mainSrc {
+		file, err := parser.ParseFile(mainFset, "main"+string(rune('0'+i))+".go", src, parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		mainFiles = append(mainFiles, file)
+	}
+	prog = llssa.NewProgram(nil)
+	mainPkg := types.NewPackage("example.com/p", "main")
+	if err := ParsePkgSyntax(prog, mainFset, mainPkg, mainFiles); err != nil {
+		t.Fatal(err)
+	}
+	if got, ok := prog.Linkname("main.demo6"); !ok || got != "main.demo" {
+		t.Fatalf("main linkname = (%q,%v), want (main.demo,true)", got, ok)
+	}
+	if got, ok := prog.Linkname("example.com/p.demo6"); !ok || got != "example.com/p.demo" {
+		t.Fatalf("raw main linkname = (%q,%v), want (example.com/p.demo,true)", got, ok)
+	}
+
+	mainCrossSrc := []string{
+		"package main\nimport _ \"unsafe\"\n//go:linkname floating main.target\n",
+		"package main\nfunc floating()\n",
+	}
+	mainCrossFset := token.NewFileSet()
+	mainCrossFiles := make([]*ast.File, 0, len(mainCrossSrc))
+	for i, src := range mainCrossSrc {
+		file, err := parser.ParseFile(mainCrossFset, "maincross"+string(rune('0'+i))+".go", src, parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		mainCrossFiles = append(mainCrossFiles, file)
+	}
+	prog = llssa.NewProgram(nil)
+	mainCrossPkg := types.NewPackage("example.com/p", "main")
+	if err := ParsePkgSyntax(prog, mainCrossFset, mainCrossPkg, mainCrossFiles); err != nil {
+		t.Fatal(err)
+	}
+	if got, ok := prog.Linkname("example.com/p.floating"); !ok || got != "example.com/p.target" {
+		t.Fatalf("raw main floating linkname = (%q,%v), want (example.com/p.target,true)", got, ok)
+	}
+
+	mainTestFset := token.NewFileSet()
+	mainTestFile, err := parser.ParseFile(mainTestFset, "main_test.go", "package main\nimport _ \"unsafe\"\n//go:linkname demo6 main.demo\nfunc demo6()\n", parser.ParseComments)
+	if err != nil {
+		t.Fatal(err)
+	}
+	prog = llssa.NewProgram(nil)
+	mainTestPkg := types.NewPackage("example.com/p", "main")
+	if err := ParsePkgSyntax(prog, mainTestFset, mainTestPkg, []*ast.File{mainTestFile}); err != nil {
+		t.Fatal(err)
+	}
+	if got, ok := prog.Linkname("example.com/p.demo6"); !ok || got != "main.demo" {
+		t.Fatalf("test-source main linkname = (%q,%v), want (main.demo,true)", got, ok)
+	}
 }
 
 func TestParsePkgSyntaxCollectsClosureEnvDirectives(t *testing.T) {

internal/build/build.go

@@ -520,11 +520,11 @@ func Build(inv Invocation) ([]Package, error) {
 		defer syntaxErrMu.Unlock()
 		return syntaxErr
 	}
-	dedup.SetPreload(func(pkg *types.Package, files []*ast.File) {
-		if llruntime.SkipToBuild(pkg.Path()) {
+	dedup.SetPreload(func(pkg *packages.Package) {
+		if llruntime.SkipToBuild(pkg.Types.Path()) {
 			return
 		}
-		if err := cl.ParsePkgSyntaxWithOptions(prog, cfg.Fset, pkg, files, preloadOptions); err != nil {
+		if err := cl.ParsePkgSyntaxWithOptions(prog, cfg.Fset, pkg.Types, pkg.Syntax, preloadOptions); err != nil {
 			recordSyntaxErr(err)
 		}
 	})

internal/packages/load.go

@@ -124,7 +124,7 @@ type aDeduper struct {
 	cache     sync.Map
 	checked   sync.Map
 	setpath   func(path string, name string) string
-	preload   func(pkg *types.Package, syntax []*ast.File)
+	preload   func(pkg *Package)
 	llgoFiles map[string][]string
 }
 
@@ -134,7 +134,7 @@ func NewDeduper() Deduper {
 	return &aDeduper{}
 }
 
-func (p Deduper) SetPreload(fn func(pkg *types.Package, syntax []*ast.File)) {
+func (p Deduper) SetPreload(fn func(pkg *Package)) {
 	p.preload = fn
 }
 
@@ -431,7 +431,7 @@ func loadPackageEx(dedup Deduper, ld *loader, lpkg *loaderPackage) {
 	})
 
 	if dedup != nil && dedup.preload != nil {
-		dedup.preload(lpkg.Types, lpkg.Syntax)
+		dedup.preload(lpkg.Package)
 	}
 
 	// type-check