diff --git a/cl/import.go b/cl/import.go index e7c2d38b40..2199e3e322 100644 --- a/cl/import.go +++ b/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 } diff --git a/cl/import_coverage_test.go b/cl/import_coverage_test.go index c2d7a2fd78..dab9f41304 100644 --- a/cl/import_coverage_test.go +++ b/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) { diff --git a/internal/build/build.go b/internal/build/build.go index 6f870695c1..72e32ecaa1 100644 --- a/internal/build/build.go +++ b/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) } }) diff --git a/internal/packages/load.go b/internal/packages/load.go index 055be834f5..7b0a49e956 100644 --- a/internal/packages/load.go +++ b/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