Skip to content

Commit b407f0e

Browse files
cpunioncpunion
committed
cl: collect closure env directives during syntax scan
1 parent 5b379e0 commit b407f0e

6 files changed

Lines changed: 174 additions & 25 deletions

File tree

‎cl/compile.go‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -618,7 +618,7 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
618618
fn := pkg.FuncOf(name)
619619
hasFreeVars := len(f.FreeVars) > 0
620620
elideFreeVarEnv := p.canElideZeroSizedClosureEnv(f)
621-
hasExplicitEnv := hasClosureEnvDirective(f)
621+
hasExplicitEnv := p.hasClosureEnvDirective(pkgTypes, f)
622622
hasCtx := hasFreeVars && !elideFreeVarEnv || hasExplicitEnv
623623
var ctx *types.Var
624624
if elideFreeVarEnv {
@@ -761,11 +761,17 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
761761
// hasClosureEnvDirective is intentionally source-only. //llgo:env may mark
762762
// only internal bodies emitted in their defining package; external env-bearing
763763
// declarations must be reconstructed explicitly with llssa.NewEnvFunc.
764-
func hasClosureEnvDirective(f *ssa.Function) bool {
764+
func (p *context) hasClosureEnvDirective(pkg *types.Package, f *ssa.Function) bool {
765765
decl, _ := f.Syntax().(*ast.FuncDecl)
766-
if decl == nil || decl.Doc == nil {
766+
if decl == nil {
767767
return false
768768
}
769+
fullName, _ := astFuncName(llssa.PathOf(pkg), decl)
770+
if enabled, ok := p.prog.ClosureEnvDirective(p.goProg.Fset, fullName, decl.Pos()); ok {
771+
return enabled
772+
}
773+
// Keep custom compilation paths that did not preload this declaration
774+
// source-compatible. Normal builds consume the ParsePkgSyntax cache above.
769775
for _, parsed := range directive.ParseGroup(decl.Doc) {
770776
if parsed.Name == "llgo:env" {
771777
return true

‎cl/import.go‎

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -276,19 +276,33 @@ func (p *context) collectSkip(line string, prefix int) {
276276
}
277277
}
278278

279-
func collectLinknameByDoc(prog llssa.Program, doc *ast.CommentGroup, fullName, inPkgName string) {
279+
// collectDeclarationDirectives caches source metadata needed after the syntax
280+
// pass. funcPos is token.NoPos for non-function declarations.
281+
func collectDeclarationDirectives(prog llssa.Program, fset *token.FileSet, doc *ast.CommentGroup, fullName, inPkgName string, funcPos token.Pos) {
280282
directives := directive.ParseGroup(doc)
283+
linkCollected := false
284+
hasClosureEnv := false
281285
for n := len(directives) - 1; n >= 0; n-- {
282-
directive := directives[n]
283-
if directive.Name != "go:linkname" && directive.Name != "llgo:link" {
284-
continue
285-
}
286-
fields := strings.Fields(directive.Args)
287-
if len(fields) >= 2 && fields[0] == inPkgName {
288-
prog.SetLinkname(fullName, strings.Join(fields[1:], " "))
289-
return
286+
item := directives[n]
287+
switch item.Name {
288+
case "go:linkname", "llgo:link":
289+
if linkCollected {
290+
continue
291+
}
292+
fields := strings.Fields(item.Args)
293+
if len(fields) >= 2 && fields[0] == inPkgName {
294+
prog.SetLinkname(fullName, strings.Join(fields[1:], " "))
295+
linkCollected = true
296+
}
297+
case "llgo:env":
298+
if funcPos.IsValid() {
299+
hasClosureEnv = true
300+
}
290301
}
291302
}
303+
if funcPos.IsValid() {
304+
prog.SetClosureEnvDirective(fset, fullName, funcPos, hasClosureEnv)
305+
}
292306
}
293307

294308
func (p *context) processLinknameByDoc(doc *ast.CommentGroup, fullName, inPkgName string, isVar, allowExport bool) bool {
@@ -563,8 +577,7 @@ const (
563577
llgoAtomicCmpXchgOK = llgoInstrBase + 0x45
564578
llgoAtomicAddReturnNew = llgoInstrBase + 0x46
565579
llgoBoolToUint8 = llgoInstrBase + 0x47
566-
// 0x48 is reserved for llgoCoroPark in the coroutine backend.
567-
llgoClosureEnv = llgoInstrBase + 0x49
580+
llgoClosureEnv = llgoInstrBase + 0x48
568581

569582
llgoAtomicOpLast = llgoAtomicOpBase + int(llssa.OpUMin)
570583
)
@@ -775,14 +788,14 @@ func ParsePkgSyntax(prog llssa.Program, fset *token.FileSet, pkg *types.Package,
775788
return err
776789
}
777790
fullName, inPkgName := astFuncName(pkgPath, decl)
778-
collectLinknameByDoc(prog, decl.Doc, fullName, inPkgName)
791+
collectDeclarationDirectives(prog, fset, decl.Doc, fullName, inPkgName, decl.Pos())
779792
ctx.processNoInterfaceByDoc(decl.Doc, fullName)
780793
case *ast.GenDecl:
781794
if decl.Tok == token.VAR {
782795
if len(decl.Specs) == 1 {
783796
if names := decl.Specs[0].(*ast.ValueSpec).Names; len(names) == 1 {
784797
inPkgName := names[0].Name
785-
collectLinknameByDoc(prog, decl.Doc, pkgPath+"."+inPkgName, inPkgName)
798+
collectDeclarationDirectives(prog, fset, decl.Doc, pkgPath+"."+inPkgName, inPkgName, token.NoPos)
786799
}
787800
}
788801
vars, err := locality.ScanPackageVar(fset, decl)

‎cl/import_coverage_test.go‎

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -251,20 +251,59 @@ func TestParsePkgSyntaxCollectsLinknames(t *testing.T) {
251251
})
252252
}
253253
prog := llssa.NewProgram(nil)
254-
collectLinknameByDoc(prog, &ast.CommentGroup{List: []*ast.Comment{{Text: "//go:linkname Other C.other"}}}, llssa.PkgRuntime+".Sigsetjmp", "Sigsetjmp")
254+
collectDeclarationDirectives(prog, nil, &ast.CommentGroup{List: []*ast.Comment{{Text: "//go:linkname Other C.other"}}}, llssa.PkgRuntime+".Sigsetjmp", "Sigsetjmp", token.NoPos)
255255
if _, ok := prog.Linkname(llssa.PkgRuntime + ".Sigsetjmp"); ok {
256256
t.Fatal("mismatched linkname was collected")
257257
}
258258
}
259259

260-
func TestCollectLinknameByDocIgnoresOtherDirectives(t *testing.T) {
260+
func TestParsePkgSyntaxCollectsClosureEnvDirectives(t *testing.T) {
261+
const src = `package p
262+
//go:linkname env C.old
263+
//llgo:env
264+
//go:linkname env C.new
265+
func env() {}
266+
267+
// llgo:env
268+
func spaced() {}
269+
270+
func plain() {}
271+
`
272+
fset := token.NewFileSet()
273+
file, err := parser.ParseFile(fset, "p.go", src, parser.ParseComments)
274+
if err != nil {
275+
t.Fatal(err)
276+
}
277+
prog := llssa.NewProgram(nil)
278+
pkg := types.NewPackage("example.com/p", "p")
279+
if err := ParsePkgSyntax(prog, fset, pkg, []*ast.File{file}); err != nil {
280+
t.Fatal(err)
281+
}
282+
if link, ok := prog.Linkname("example.com/p.env"); !ok || link != "C.new" {
283+
t.Fatalf("combined declaration linkname = (%q, %v), want (C.new, true)", link, ok)
284+
}
285+
want := map[string]bool{"env": true, "spaced": true, "plain": false}
286+
for _, node := range file.Decls {
287+
decl := node.(*ast.FuncDecl)
288+
fullName, _ := astFuncName(pkg.Path(), decl)
289+
got, ok := prog.ClosureEnvDirective(fset, fullName, decl.Pos())
290+
if !ok || got != want[decl.Name.Name] {
291+
t.Fatalf("ClosureEnvDirective(%s) = (%v, %v), want (%v, true)", decl.Name.Name, got, ok, want[decl.Name.Name])
292+
}
293+
}
294+
if _, ok := prog.ClosureEnvDirective(fset, "example.com/p.missing", token.NoPos); ok {
295+
t.Fatal("missing declaration unexpectedly has cached directives")
296+
}
297+
}
298+
299+
func TestCollectDeclarationDirectivesIgnoresOtherDirectives(t *testing.T) {
261300
prog := llssa.NewProgram(nil)
262301
doc := &ast.CommentGroup{List: []*ast.Comment{
263302
{Text: "//go:noinline"},
264303
{Text: "//llgo:tls"},
265304
}}
266305
const fullName = "example.com/p.Value"
267-
collectLinknameByDoc(prog, doc, fullName, "Value")
306+
collectDeclarationDirectives(prog, nil, doc, fullName, "Value", token.NoPos)
268307
if _, ok := prog.Linkname(fullName); ok {
269308
t.Fatal("non-link directives installed a linkname")
270309
}

‎cl/rewrite_internal_test.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,39 @@ func compileWithRewritesTarget(t *testing.T, src string, rewrites map[string]str
5757
return ret.String()
5858
}
5959

60+
func TestClosureEnvDirectiveFallbackWithoutSyntaxPreload(t *testing.T) {
61+
const src = `package fallback
62+
//llgo:env
63+
func withEnv() {}
64+
func plain() {}
65+
`
66+
fset := token.NewFileSet()
67+
file, err := parser.ParseFile(fset, "fallback.go", src, parser.ParseComments)
68+
if err != nil {
69+
t.Fatal(err)
70+
}
71+
importer := gpackages.NewImporter(fset)
72+
pkg, _, err := ssautil.BuildPackage(&types.Config{Importer: importer}, fset,
73+
types.NewPackage("fallback", "fallback"), []*ast.File{file}, ssa.SanityCheckFunctions)
74+
if err != nil {
75+
t.Fatal(err)
76+
}
77+
prog := ssatest.NewProgramEx(t, nil, importer)
78+
ctx := &context{prog: prog, goProg: pkg.Prog}
79+
for _, test := range []struct {
80+
name string
81+
want bool
82+
}{
83+
{"withEnv", true},
84+
{"plain", false},
85+
} {
86+
fn := pkg.Func(test.name)
87+
if got := ctx.hasClosureEnvDirective(pkg.Pkg, fn); got != test.want {
88+
t.Fatalf("hasClosureEnvDirective(%s) = %v, want %v", test.name, got, test.want)
89+
}
90+
}
91+
}
92+
6093
func TestClosureEnvIntrinsicRequiresEnvBearingEntry(t *testing.T) {
6194
valid := `package closureenv
6295

‎ssa/closure_env_test.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,38 @@ import (
1212
"github.com/xgo-dev/llvm"
1313
)
1414

15+
func TestClosureEnvDirectiveCacheUsesSourceIdentity(t *testing.T) {
16+
prog := NewProgram(nil)
17+
defer prog.Dispose()
18+
fset := token.NewFileSet()
19+
otherFset := token.NewFileSet()
20+
const (
21+
name = "example.com/p.entry"
22+
pos = token.Pos(7)
23+
)
24+
prog.SetClosureEnvDirective(fset, name, pos, true)
25+
if enabled, ok := prog.ClosureEnvDirective(fset, name, pos); !ok || !enabled {
26+
t.Fatalf("ClosureEnvDirective() = (%v, %v), want (true, true)", enabled, ok)
27+
}
28+
for _, key := range []struct {
29+
fset *token.FileSet
30+
name string
31+
pos token.Pos
32+
}{
33+
{otherFset, name, pos},
34+
{fset, "example.com/p.alias", pos},
35+
{fset, name, pos + 1},
36+
} {
37+
if _, ok := prog.ClosureEnvDirective(key.fset, key.name, key.pos); ok {
38+
t.Fatalf("distinct source declaration (%p, %q, %d) shared cached directives", key.fset, key.name, key.pos)
39+
}
40+
}
41+
prog.SetClosureEnvDirective(fset, name, pos, false)
42+
if enabled, ok := prog.ClosureEnvDirective(fset, name, pos); !ok || enabled {
43+
t.Fatalf("updated ClosureEnvDirective() = (%v, %v), want (false, true)", enabled, ok)
44+
}
45+
}
46+
1547
func TestClosureEnvABIForTarget(t *testing.T) {
1648
tests := []struct {
1749
triple string

‎ssa/package.go‎

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -221,12 +221,13 @@ type aProgram struct {
221221

222222
printfTy *types.Signature
223223

224-
paramObjPtr_ *types.Var
225-
linknameMu sync.RWMutex
226-
linkname map[string]string // pkgPath.nameInPkg => linkname
227-
localities *localityInfos
228-
noInterface map[string]none // pkgPath.T.method or pkgPath.(*T).method
229-
abiSymbol map[string]*AbiSymbol // abi symbol name => AbiSymbol
224+
paramObjPtr_ *types.Var
225+
linknameMu sync.RWMutex
226+
linkname map[string]string // pkgPath.nameInPkg => linkname
227+
closureEnvDirectives sync.Map // closureEnvDirectiveKey => bool
228+
localities *localityInfos
229+
noInterface map[string]none // pkgPath.T.method or pkgPath.(*T).method
230+
abiSymbol map[string]*AbiSymbol // abi symbol name => AbiSymbol
230231

231232
ptrSize int
232233

@@ -425,6 +426,31 @@ func (p Program) Linkname(name string) (link string, ok bool) {
425426
return
426427
}
427428

429+
type closureEnvDirectiveKey struct {
430+
fset *token.FileSet
431+
name string
432+
pos token.Pos
433+
}
434+
435+
// SetClosureEnvDirective records whether a source function declaration has the
436+
// llgo:env directive. name and pos identify the source declaration rather
437+
// than its resolved linker symbol, so aliases retain independent ABI metadata.
438+
func (p Program) SetClosureEnvDirective(fset *token.FileSet, name string, pos token.Pos, enabled bool) {
439+
key := closureEnvDirectiveKey{fset: fset, name: name, pos: pos}
440+
p.closureEnvDirectives.Store(key, enabled)
441+
}
442+
443+
// ClosureEnvDirective reports the cached llgo:env state for a source
444+
// function declaration. ok is false when that declaration was not scanned.
445+
func (p Program) ClosureEnvDirective(fset *token.FileSet, name string, pos token.Pos) (enabled, ok bool) {
446+
key := closureEnvDirectiveKey{fset: fset, name: name, pos: pos}
447+
value, ok := p.closureEnvDirectives.Load(key)
448+
if ok {
449+
enabled = value.(bool)
450+
}
451+
return enabled, ok
452+
}
453+
428454
func (p Program) runtime() *types.Package {
429455
if p.rt == nil {
430456
p.rt = p.rtget()

0 commit comments

Comments
 (0)