ssa,cl,build: release per-compile memory in in-process drivers

cl/caller_frame_test.go

@@ -196,25 +196,26 @@ func callWorker(w workerIface) { w.Work() }
 func workerHot() { var w workerIface = workerImpl{}; callWorker(w) }
 func plain() {}
 `)
-	if !packageUsesRuntimeCaller(ssapkg) {
+	callerCaches := NewCallerTracking()
+	if !packageUsesRuntimeCaller(callerCaches, ssapkg) {
 		t.Fatal("package should report runtime caller usage")
 	}
-	if !fnUsesRuntimeCaller(ssapkg.Func("direct")) {
+	if !fnUsesRuntimeCaller(callerCaches, ssapkg.Func("direct")) {
 		t.Fatal("direct runtime.Caller use should be detected")
 	}
-	if !fnUsesRuntimeCaller(ssapkg.Func("indirect")) {
+	if !fnUsesRuntimeCaller(callerCaches, ssapkg.Func("indirect")) {
 		t.Fatal("transitive runtime.Caller use should be detected")
 	}
-	if !fnUsesRuntimeCaller(ssapkg.Func("stack")) {
+	if !fnUsesRuntimeCaller(callerCaches, ssapkg.Func("stack")) {
 		t.Fatal("runtime/debug.Stack use should be detected")
 	}
-	if !fnUsesRuntimeCaller(ssapkg.Func("anonOnly")) {
+	if !fnUsesRuntimeCaller(callerCaches, ssapkg.Func("anonOnly")) {
 		t.Fatal("runtime caller use in anonymous functions should be detected")
 	}
-	if fnUsesRuntimeCaller(ssapkg.Func("plain")) {
+	if fnUsesRuntimeCaller(callerCaches, ssapkg.Func("plain")) {
 		t.Fatal("plain function should not report runtime caller usage")
 	}
-	runtimeCallerFuncs := runtimeCallerFuncSet(ssapkg)
+	runtimeCallerFuncs := runtimeCallerFuncSet(callerCaches, ssapkg)
 	for _, name := range []string{"dynamic", "dynamicCaller", "interfaceDispatch", "interfaceCaller", "closureLayer", "closureCaller"} {
 		if !runtimeCallerFuncs[ssapkg.Func(name)] {
 			t.Fatalf("%s should be tracked because dynamic calls may reach runtime stack APIs", name)
@@ -266,13 +267,14 @@ func FuncForPC(pc uintptr) uintptr { return 0 }
 }
 
 func TestRuntimeCallerAnalysisEdgeCases(t *testing.T) {
-	if fnUsesRuntimeCaller(nil) {
+	callerCaches := NewCallerTracking()
+	if fnUsesRuntimeCaller(callerCaches, nil) {
 		t.Fatal("nil function should not use runtime caller metadata")
 	}
-	if fnUsesRuntimeCaller(&gossa.Function{}) {
+	if fnUsesRuntimeCaller(callerCaches, &gossa.Function{}) {
 		t.Fatal("function without a package should not use runtime caller metadata")
 	}
-	if runtimeCallerFuncSet(nil) != nil {
+	if runtimeCallerFuncSet(callerCaches, nil) != nil {
 		t.Fatal("nil package should have no runtime caller set")
 	}
 	if fnHasDirectRuntimeCaller(nil) {
@@ -381,7 +383,7 @@ type T struct{}
 func (T) Call() { runtime.Caller(0) }
 var _ = T{}
 `)
-	methodOnlySet := runtimeCallerFuncSet(methodOnlyPkg)
+	methodOnlySet := runtimeCallerFuncSet(NewCallerTracking(), methodOnlyPkg)
 	if methodOnlySet == nil {
 		t.Fatal("a method calling runtime.Caller must be tracked (slog.(*Logger).Info escaped exactly this way)")
 	}
@@ -445,7 +447,7 @@ func f() { runtime.Caller(0) }
 				fn:                 fn,
 				goFn:               goFn,
 				trackCallerFrames:  tt.track,
-				runtimeCallerFuncs: runtimeCallerFuncSet(ssapkg),
+				runtimeCallerFuncs: runtimeCallerFuncSet(NewCallerTracking(), ssapkg),
 			}
 			if got := ctx.shouldTrackCallerFrames(); got != tt.want {
 				t.Fatalf("shouldTrackCallerFrames() = %v, want %v", got, tt.want)
@@ -650,7 +652,7 @@ func top() {
 		goFn:               ssapkg.Func("top"),
 		fset:               token.NewFileSet(),
 		trackCallerFrames:  true,
-		runtimeCallerFuncs: runtimeCallerFuncSet(ssapkg),
+		runtimeCallerFuncs: runtimeCallerFuncSet(NewCallerTracking(), ssapkg),
 	}
 	var b llssa.Builder
 	ctx.pushCallerLocationFrame(b, nil)
@@ -940,10 +942,11 @@ func Logs() bool { return dep.Where() }
 
 func Plain() int { return dep.Quiet() }
 `)
-	if !runtimeCallerBaseSet(depSSA)[depSSA.Func("Where")] {
+	crossCaches := NewCallerTracking()
+	if !runtimeCallerBaseSet(crossCaches, depSSA)[depSSA.Func("Where")] {
 		t.Fatal("dep.Where must be in its own package's base set")
 	}
-	set := runtimeCallerFuncSet(rootSSA)
+	set := runtimeCallerFuncSet(crossCaches, rootSSA)
 	if !set[rootSSA.Func("Logs")] {
 		t.Fatal("caller of a pc-consuming function must be tracked (criterion 2)")
 	}
@@ -966,7 +969,7 @@ func pinned() {}
 
 func helper() {}
 `)
-	set := runtimeCallerFuncSet(ssapkg)
+	set := runtimeCallerFuncSet(NewCallerTracking(), ssapkg)
 	for _, name := range []string{"main", "init", "pinned"} {
 		if !set[ssapkg.Func(name)] {
 			t.Fatalf("%s must be pinned in the tracking set", name)

cl/caller_tracking_compile_test.go

@@ -0,0 +1,58 @@
+//go:build !llgo
+// +build !llgo
+
+package cl_test
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"go/types"
+	"runtime"
+	"testing"
+
+	"github.com/goplus/gogen/packages"
+	"github.com/goplus/llgo/cl"
+	"github.com/goplus/llgo/ssa/ssatest"
+	"golang.org/x/tools/go/ssa"
+	"golang.org/x/tools/go/ssa/ssautil"
+)
+
+// TestNewPackageExWithEmbedSharedTracking compiles with a caller-provided
+// CallerTracking, the multi-package driver pattern where one instance is
+// shared across every package of a compilation (like patches).
+func TestNewPackageExWithEmbedSharedTracking(t *testing.T) {
+	src := `package foo
+
+import "runtime"
+
+func Where() string {
+	_, file, _, _ := runtime.Caller(0)
+	return file
+}
+`
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "foo.go", src, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("ParseFile failed: %v", err)
+	}
+	files := []*ast.File{f}
+
+	imp := packages.NewImporter(fset)
+	mode := ssa.SanityCheckFunctions | ssa.InstantiateGenerics
+	fooPkg, _, err := ssautil.BuildPackage(&types.Config{Importer: imp}, fset, types.NewPackage(f.Name.Name, f.Name.Name), files, mode)
+	if err != nil {
+		t.Fatalf("BuildPackage failed: %v", err)
+	}
+
+	prog := ssatest.NewProgramEx(t, nil, imp)
+	prog.TypeSizes(types.SizesFor("gc", runtime.GOARCH))
+	ct := cl.NewCallerTracking()
+	ret, _, err := cl.NewPackageExWithEmbed(prog, ct, nil, nil, fooPkg, files, nil)
+	if err != nil {
+		t.Fatalf("NewPackageExWithEmbed failed: %v", err)
+	}
+	if ret.String() == "" {
+		t.Fatal("empty IR")
+	}
+}

cl/compile.go

@@ -1881,6 +1881,10 @@ func NewPackage(prog llssa.Program, pkg *ssa.Package, files []*ast.File) (ret ll
 	return
 }
 
+// NewPackageEx and NewPackage compile as a one-shot compilation: each
+// call gets fresh caller-tracking memoization. Multi-package drivers
+// use NewPackageExWithEmbed with a shared CallerTracking instead.
+
 // NewPackageEx compiles a Go package to LLVM IR package.
 //
 // Parameters:
@@ -1893,17 +1897,21 @@ func NewPackage(prog llssa.Program, pkg *ssa.Package, files []*ast.File) (ret ll
 // The rewrites map uses short variable names (without package qualifier) and
 // only affects string-typed globals defined in the current package.
 func NewPackageEx(prog llssa.Program, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File) (ret llssa.Package, externs []string, err error) {
-	return newPackageEx(prog, patches, rewrites, pkg, files, nil)
+	return newPackageEx(prog, nil, patches, rewrites, pkg, files, nil)
 }
 
 // NewPackageExWithEmbed compiles a package using pre-loaded go:embed metadata.
 //
 // This avoids re-scanning directives when the caller already loaded them.
-func NewPackageExWithEmbed(prog llssa.Program, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File, embedMap goembed.VarMap) (ret llssa.Package, externs []string, err error) {
-	return newPackageEx(prog, patches, rewrites, pkg, files, &embedMap)
+// ct carries the compilation-scoped caller-tracking memoization; drivers
+// compiling multiple packages pass the same instance for every package
+// of one compilation (like patches). nil means one-shot: a fresh
+// instance is created for this call.
+func NewPackageExWithEmbed(prog llssa.Program, ct *CallerTracking, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File, embedMap goembed.VarMap) (ret llssa.Package, externs []string, err error) {
+	return newPackageEx(prog, ct, patches, rewrites, pkg, files, &embedMap)
 }
 
-func newPackageEx(prog llssa.Program, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File, embedMap *goembed.VarMap) (ret llssa.Package, externs []string, err error) {
+func newPackageEx(prog llssa.Program, ct *CallerTracking, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File, embedMap *goembed.VarMap) (ret llssa.Package, externs []string, err error) {
 	pkgProg := pkg.Prog
 	pkgTypes := pkg.Pkg
 	oldTypes := pkgTypes
@@ -1922,6 +1930,9 @@ func newPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 		ret.InitDebug(pkgName, pkgPath, pkgProg.Fset)
 	}
 
+	if ct == nil {
+		ct = NewCallerTracking()
+	}
 	ctx := &context{
 		prog:             prog,
 		pkg:              ret,
@@ -1941,8 +1952,8 @@ func newPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 		cgoSymbols: make([]string, 0, 128),
 		rewrites:   rewrites,
 
-		trackCallerFrames:  filesUseRuntimeCaller(files) || packageUsesRuntimeCaller(pkg),
-		runtimeCallerFuncs: runtimeCallerFuncSet(pkg),
+		trackCallerFrames:  filesUseRuntimeCaller(files) || packageUsesRuntimeCaller(ct, pkg),
+		runtimeCallerFuncs: runtimeCallerFuncSet(ct, pkg),
 	}
 	if embedMap != nil {
 		ctx.embedMap = *embedMap

cl/embed_with_map_compile_test.go

@@ -57,7 +57,7 @@ var files embed.FS
 
 	prog := ssatest.NewProgramEx(t, nil, imp)
 	prog.TypeSizes(types.SizesFor("gc", runtime.GOARCH))
-	ret, _, err := cl.NewPackageExWithEmbed(prog, nil, nil, fooPkg, files, embedMap)
+	ret, _, err := cl.NewPackageExWithEmbed(prog, nil, nil, nil, fooPkg, files, embedMap)
 	if err != nil {
 		t.Fatalf("NewPackageExWithEmbed failed: %v", err)
 	}

cl/instr.go

@@ -26,7 +26,6 @@ import (
 	"os"
 	"regexp"
 	"strings"
-	"sync"
 
 	"golang.org/x/tools/go/ssa"
 
@@ -880,18 +879,18 @@ func canTrackCallerFramesForPackage(pkgPath string) bool {
 		!strings.HasPrefix(pkgPath, "github.com/goplus/llgo/runtime/internal/")
 }
 
-func packageUsesRuntimeCaller(pkg *ssa.Package) bool {
-	return len(runtimeCallerFuncSet(pkg)) != 0
+func packageUsesRuntimeCaller(c *CallerTracking, pkg *ssa.Package) bool {
+	return len(runtimeCallerFuncSet(c, pkg)) != 0
 }
 
-func fnUsesRuntimeCaller(fn *ssa.Function) bool {
+func fnUsesRuntimeCaller(c *CallerTracking, fn *ssa.Function) bool {
 	if fn == nil {
 		return false
 	}
 	if fn.Pkg == nil {
 		return fnHasDirectRuntimeCaller(fn)
 	}
-	return runtimeCallerFuncSet(fn.Pkg)[fn]
+	return runtimeCallerFuncSet(c, fn.Pkg)[fn]
 }
 
 // runtimeCallerFuncSet is the per-package tracking set: functions that
@@ -910,15 +909,14 @@ func fnUsesRuntimeCaller(fn *ssa.Function) bool {
 // (criterion 1 alone), so tracking extends exactly one call level past a
 // pc-consuming package and does not cascade through arbitrary wrapper
 // layers; multi-package wrapper chains remain the P4 inline-tree's job.
-func runtimeCallerFuncSet(pkg *ssa.Package) map[*ssa.Function]bool {
+func runtimeCallerFuncSet(c *CallerTracking, pkg *ssa.Package) map[*ssa.Function]bool {
 	if pkg == nil {
 		return nil
 	}
-	if v, ok := runtimeCallerExtendedCache.Load(pkg); ok {
-		set, _ := v.(map[*ssa.Function]bool)
+	if set, ok := c.extended[pkg]; ok {
 		return set
 	}
-	base := runtimeCallerBaseSet(pkg)
+	base := runtimeCallerBaseSet(c, pkg)
 	out := make(map[*ssa.Function]bool, len(base))
 	for fn := range base {
 		out[fn] = true
@@ -952,22 +950,41 @@ func runtimeCallerFuncSet(pkg *ssa.Package) map[*ssa.Function]bool {
 			if !canTrackCallerFramesForPackage(callee.Pkg.Pkg.Path()) {
 				return
 			}
-			if runtimeCallerBaseSet(callee.Pkg)[callee] {
+			if runtimeCallerBaseSet(c, callee.Pkg)[callee] {
 				out[fn] = true
 			}
 		})
 	}
 	if len(out) == 0 {
 		out = nil
 	}
-	runtimeCallerExtendedCache.Store(pkg, out)
+	c.extended[pkg] = out
 	return out
 }
 
-var (
-	runtimeCallerBaseCache     sync.Map // *ssa.Package -> map[*ssa.Function]bool
-	runtimeCallerExtendedCache sync.Map // *ssa.Package -> map[*ssa.Function]bool
-)
+// CallerTracking memoizes the per-package caller-tracking sets for one
+// compilation. Like Patches, it is compilation-scoped state owned by the
+// driver: create one per compilation and pass it to every
+// NewPackageExWithEmbed call of that compilation, so cross-package
+// queries (criterion 2 below) hit the memoization. It must not outlive
+// the compilation — the maps are keyed by *ssa.Package with
+// *ssa.Function values, so anything longer-lived would pin every
+// compiled package's go/types and go/ssa graphs. Plain maps are enough:
+// packages of one compilation are compiled sequentially (the LLVM
+// context is not thread-safe).
+type CallerTracking struct {
+	base     map[*ssa.Package]map[*ssa.Function]bool
+	extended map[*ssa.Package]map[*ssa.Function]bool
+}
+
+// NewCallerTracking creates the caller-tracking memoization for one
+// compilation.
+func NewCallerTracking() *CallerTracking {
+	return &CallerTracking{
+		base:     make(map[*ssa.Package]map[*ssa.Function]bool),
+		extended: make(map[*ssa.Package]map[*ssa.Function]bool),
+	}
+}
 
 func isProgramUniqueFrame(pkg *ssa.Package, fn *ssa.Function) bool {
 	if fn == nil || fn.Parent() != nil {
@@ -980,16 +997,15 @@ func isProgramUniqueFrame(pkg *ssa.Package, fn *ssa.Function) bool {
 	return name == "main" && pkg.Pkg != nil && pkg.Pkg.Name() == "main"
 }
 
-func runtimeCallerBaseSet(pkg *ssa.Package) map[*ssa.Function]bool {
+func runtimeCallerBaseSet(c *CallerTracking, pkg *ssa.Package) map[*ssa.Function]bool {
 	if pkg == nil {
 		return nil
 	}
-	if v, ok := runtimeCallerBaseCache.Load(pkg); ok {
-		set, _ := v.(map[*ssa.Function]bool)
+	if set, ok := c.base[pkg]; ok {
 		return set
 	}
 	set := computeRuntimeCallerBaseSet(pkg)
-	runtimeCallerBaseCache.Store(pkg, set)
+	c.base[pkg] = set
 	return set
 }
 

internal/build/build.go

@@ -322,6 +322,15 @@ func Do(args []string, conf *Config) ([]Package, error) {
 	}
 
 	prog := llssa.NewProgram(target)
+	if conf.Mode != ModeGen {
+		// ModeGen callers (llgen and the golden suites) read LPkg.String()
+		// after Do returns and dispose the program themselves; every other
+		// mode's outputs are files or a spawned process, so the compile's
+		// LLVM context can be released when Do finishes. In-process
+		// drivers that build many packages per process (the cltest run
+		// harness) otherwise accumulate every compile's C++-side memory.
+		defer prog.Dispose()
+	}
 	prog.EnableGoGlobalDCE(conf.goGlobalDCEEnabled())
 	if conf.PthreadStackSize > 0 {
 		prog.SetPthreadStackSize(uint64(conf.PthreadStackSize))
@@ -432,7 +441,8 @@ func Do(args []string, conf *Config) ([]Package, error) {
 
 	output := conf.OutFile != ""
 	ctx := &context{env: env, conf: cfg, progSSA: progSSA, prog: prog, dedup: dedup,
-		patches: patches, built: make(map[string]none), initial: initial, mode: mode,
+		patches: patches, callerTracking: cl.NewCallerTracking(),
+		built: make(map[string]none), initial: initial, mode: mode,
 		fingerprinting: make(map[string]bool),
 		pkgs:           map[*packages.Package]Package{},
 		pkgByID:        map[string]Package{},
@@ -609,6 +619,7 @@ type context struct {
 	prog           llssa.Program
 	dedup          packages.Deduper
 	patches        cl.Patches
+	callerTracking *cl.CallerTracking
 	built          map[string]none
 	fingerprinting map[string]bool
 	initial        []*packages.Package
@@ -1370,7 +1381,7 @@ func buildPkg(ctx *context, aPkg *aPackage, verbose bool) error {
 		return fmt.Errorf("load go:embed directives for %s failed: %w", pkgPath, err)
 	}
 
-	ret, externs, err := cl.NewPackageExWithEmbed(ctx.prog, ctx.patches, aPkg.rewriteVars, aPkg.SSA, syntax, embedMap)
+	ret, externs, err := cl.NewPackageExWithEmbed(ctx.prog, ctx.callerTracking, ctx.patches, aPkg.rewriteVars, aPkg.SSA, syntax, embedMap)
 	check(err)
 
 	aPkg.LPkg = ret