runtime: add memprofile attribution support

cl/compile.go

@@ -157,25 +157,27 @@ type pkgInfo struct {
 type none = struct{}
 
 type context struct {
-	prog                 llssa.Program
-	pkg                  llssa.Package
-	fn                   llssa.Function
-	goFn                 *ssa.Function
-	fset                 *token.FileSet
-	goProg               *ssa.Program
-	goTyps               *types.Package
-	goPkg                *ssa.Package
-	pyMod                string
-	skips                map[string]none
-	loaded               map[*types.Package]*pkgInfo // loaded packages
-	bvals                map[ssa.Value]llssa.Expr    // block values
-	methodNilDerefChecks map[*ssa.UnOp]none
-	vargs                map[*ssa.Alloc][]llssa.Expr // varargs
-	funcs                map[*ssa.Function]llssa.Function
-	linkOnceFns          map[*ssa.Function]none
-	stackDefers          map[*ssa.Function]bool
-	anonDefers           map[*ssa.Function]bool
-	paramDIVars          map[*types.Var]llssa.DIVar
+	prog                  llssa.Program
+	pkg                   llssa.Package
+	fn                    llssa.Function
+	goFn                  *ssa.Function
+	fset                  *token.FileSet
+	goProg                *ssa.Program
+	goTyps                *types.Package
+	goPkg                 *ssa.Package
+	pyMod                 string
+	skips                 map[string]none
+	loaded                map[*types.Package]*pkgInfo // loaded packages
+	bvals                 map[ssa.Value]llssa.Expr    // block values
+	methodNilDerefChecks  map[*ssa.UnOp]none
+	vargs                 map[*ssa.Alloc][]llssa.Expr // varargs
+	funcs                 map[*ssa.Function]llssa.Function
+	linkOnceFns           map[*ssa.Function]none
+	stackDefers           map[*ssa.Function]bool
+	anonDefers            map[*ssa.Function]bool
+	paramDIVars           map[*types.Var]llssa.DIVar
+	noInlineForMemProfile bool
+	memProfileInstrument  bool
 
 	patches          Patches
 	blkInfos         []blocks.Info
@@ -432,6 +434,23 @@ func hasInstantiatedRecv(recv *types.Var) bool {
 	return false
 }
 
+func (p *context) applyNoInline(fn llssa.Function) {
+	if disableInline || p.noInlineForMemProfile {
+		fn.Inline(llssa.NoInline)
+	}
+	if p.noInlineForMemProfile {
+		fn.DisableTailCalls()
+	}
+}
+
+func memProfileFunctionName(name string) string {
+	const commandLineArguments = "command-line-arguments."
+	if strings.HasPrefix(name, commandLineArguments) {
+		return "main." + name[len(commandLineArguments):]
+	}
+	return name
+}
+
 func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Function, llssa.PyObjRef, int) {
 	pkgTypes, name, ftype := p.funcName(f)
 	if ftype != goFunc {
@@ -469,10 +488,8 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 	}
 	if fn == nil {
 		fn = pkg.NewFuncEx(name, sig, llssa.Background(ftype), hasCtx, p.needsLinkOnce(f))
-		if disableInline {
-			fn.Inline(llssa.NoInline)
-		}
 	}
+	p.applyNoInline(fn)
 	p.funcs[f] = fn
 	isCgo := isCgoExternSymbol(f)
 	if nblk := len(f.Blocks); nblk > 0 {
@@ -499,13 +516,15 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 		}
 		dbgEnabled := enableDbg && (f == nil || f.Origin() == nil)
 		dbgSymsEnabled := enableDbgSyms && (f == nil || f.Origin() == nil)
+		instrumentMemProfile := p.noInlineForMemProfile && !isCgo
 		p.inits = append(p.inits, func() {
-			oldFn, oldGoFn, oldMethodNilDerefChecks := p.fn, p.goFn, p.methodNilDerefChecks
+			oldFn, oldGoFn, oldMethodNilDerefChecks, oldMemProfileInstrument := p.fn, p.goFn, p.methodNilDerefChecks, p.memProfileInstrument
 			p.fn = fn
 			p.goFn = f
+			p.memProfileInstrument = instrumentMemProfile
 			p.state = state // restore pkgState when compiling funcBody
 			defer func() {
-				p.fn, p.goFn, p.methodNilDerefChecks = oldFn, oldGoFn, oldMethodNilDerefChecks
+				p.fn, p.goFn, p.methodNilDerefChecks, p.memProfileInstrument = oldFn, oldGoFn, oldMethodNilDerefChecks, oldMemProfileInstrument
 			}()
 			p.phis = nil
 			if dbgSymsEnabled {
@@ -627,6 +646,9 @@ func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, do
 	var instrs = block.Instrs[n:]
 	var ret = fn.Block(block.Index)
 	b.SetBlock(ret)
+	if block.Index == 0 && p.memProfileInstrument {
+		b.MemProfileEnter(memProfileFunctionName(fn.Name()))
+	}
 	if block.Index == 0 && enableCallTracing && !strings.HasPrefix(fn.Name(), "github.com/goplus/llgo/runtime/internal/runtime.Print") {
 		b.Printf("call " + fn.Name() + "\n\x00")
 	}
@@ -1383,6 +1405,9 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 		if p.returnNeedsImplicitRunDefers(v) {
 			b.RunDefers()
 		}
+		if p.memProfileInstrument {
+			b.MemProfileExit()
+		}
 		b.Return(results...)
 	case *ssa.If:
 		fn := p.fn
@@ -1683,6 +1708,52 @@ func NewPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 	return newPackageEx(prog, patches, rewrites, pkg, files, nil)
 }
 
+func packageUsesRuntimeMemProfile(files []*ast.File) bool {
+	for _, file := range files {
+		runtimeNames := make(map[string]none)
+		for _, imp := range file.Imports {
+			path, err := strconv.Unquote(imp.Path.Value)
+			if err != nil || path != "runtime" {
+				continue
+			}
+			if imp.Name != nil {
+				if imp.Name.Name == "_" || imp.Name.Name == "." {
+					continue
+				}
+				runtimeNames[imp.Name.Name] = none{}
+				continue
+			}
+			runtimeNames["runtime"] = none{}
+		}
+		if len(runtimeNames) == 0 {
+			continue
+		}
+		found := false
+		ast.Inspect(file, func(n ast.Node) bool {
+			sel, ok := n.(*ast.SelectorExpr)
+			if !ok {
+				return true
+			}
+			if sel.Sel.Name != "MemProfile" && sel.Sel.Name != "MemProfileRate" {
+				return true
+			}
+			x, ok := sel.X.(*ast.Ident)
+			if !ok {
+				return true
+			}
+			if _, ok := runtimeNames[x.Name]; !ok {
+				return true
+			}
+			found = true
+			return false
+		})
+		if found {
+			return true
+		}
+	}
+	return false
+}
+
 // NewPackageExWithEmbed compiles a package using pre-loaded go:embed metadata.
 //
 // This avoids re-scanning directives when the caller already loaded them.
@@ -1710,18 +1781,19 @@ func newPackageEx(prog llssa.Program, patches Patches, rewrites map[string]strin
 	}
 
 	ctx := &context{
-		prog:             prog,
-		pkg:              ret,
-		fset:             pkgProg.Fset,
-		goProg:           pkgProg,
-		goTyps:           pkgTypes,
-		goPkg:            pkg,
-		patches:          patches,
-		skips:            make(map[string]none),
-		vargs:            make(map[*ssa.Alloc][]llssa.Expr),
-		funcs:            make(map[*ssa.Function]llssa.Function),
-		linkOnceFns:      make(map[*ssa.Function]none),
-		addrOfFieldAddrs: collectAddrOfFieldSelectors(files),
+		prog:                  prog,
+		pkg:                   ret,
+		fset:                  pkgProg.Fset,
+		goProg:                pkgProg,
+		goTyps:                pkgTypes,
+		goPkg:                 pkg,
+		patches:               patches,
+		noInlineForMemProfile: packageUsesRuntimeMemProfile(files),
+		skips:                 make(map[string]none),
+		vargs:                 make(map[*ssa.Alloc][]llssa.Expr),
+		funcs:                 make(map[*ssa.Function]llssa.Function),
+		linkOnceFns:           make(map[*ssa.Function]none),
+		addrOfFieldAddrs:      collectAddrOfFieldSelectors(files),
 		loaded: map[*types.Package]*pkgInfo{
 			types.Unsafe: {kind: PkgDeclOnly}, // TODO(xsw): PkgNoInit or PkgDeclOnly?
 		},

cl/instr.go

@@ -666,9 +666,7 @@ func (p *context) funcOf(fn *ssa.Function) (aFn llssa.Function, pyFn llssa.PyObj
 			}
 			sig := p.patchType(fn.Signature).(*types.Signature)
 			aFn = pkg.NewFuncEx(name, sig, llssa.Background(ftype), false, p.needsLinkOnce(fn))
-			if disableInline {
-				aFn.Inline(llssa.NoInline)
-			}
+			p.applyNoInline(aFn)
 		}
 	}
 	return

cl/memprofile_instrumentation_test.go

@@ -0,0 +1,170 @@
+//go:build !llgo
+// +build !llgo
+
+/*
+ * Copyright (c) 2026 The XGo Authors (xgo.dev). All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cl
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"strings"
+	"testing"
+)
+
+func TestPackageUsesRuntimeMemProfile(t *testing.T) {
+	tests := []struct {
+		name string
+		src  string
+		want bool
+	}{
+		{
+			name: "no runtime import",
+			src: `package foo
+func f() {}
+`,
+		},
+		{
+			name: "runtime mem profile call",
+			src: `package foo
+import "runtime"
+func f() {
+	runtime.MemProfile(nil, false)
+}
+`,
+			want: true,
+		},
+		{
+			name: "runtime mem profile rate",
+			src: `package foo
+import "runtime"
+var _ = runtime.MemProfileRate
+`,
+			want: true,
+		},
+		{
+			name: "renamed runtime import",
+			src: `package foo
+import rt "runtime"
+var _ = rt.MemProfileRate
+`,
+			want: true,
+		},
+		{
+			name: "blank runtime import",
+			src: `package foo
+import _ "runtime"
+func f() {}
+`,
+		},
+		{
+			name: "dot runtime import",
+			src: `package foo
+import . "runtime"
+var _ = MemProfileRate
+`,
+		},
+		{
+			name: "other runtime selector",
+			src: `package foo
+import "runtime"
+var _ = runtime.GOOS
+`,
+		},
+		{
+			name: "selector on non runtime value",
+			src: `package foo
+import "runtime"
+var _ = runtime.GOOS
+type profiler struct{}
+var p profiler
+var _ = p.MemProfile
+`,
+		},
+		{
+			name: "selector base is not identifier",
+			src: `package foo
+import "runtime"
+type profiler struct{}
+var p profiler
+var _ = (p).MemProfile
+`,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fset := token.NewFileSet()
+			file, err := parser.ParseFile(fset, "foo.go", tt.src, parser.ParseComments)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got := packageUsesRuntimeMemProfile([]*ast.File{file}); got != tt.want {
+				t.Fatalf("packageUsesRuntimeMemProfile() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+
+	badImport := &ast.File{
+		Imports: []*ast.ImportSpec{{
+			Path: &ast.BasicLit{Kind: token.STRING, Value: "runtime"},
+		}},
+	}
+	if packageUsesRuntimeMemProfile([]*ast.File{badImport}) {
+		t.Fatal("bad import literal should not enable memprofile instrumentation")
+	}
+}
+
+func TestMemProfileFunctionName(t *testing.T) {
+	tests := []struct {
+		name string
+		want string
+	}{
+		{name: "command-line-arguments.main", want: "main.main"},
+		{name: "command-line-arguments.profiledAlloc", want: "main.profiledAlloc"},
+		{name: "example.com/mod.profiledAlloc", want: "example.com/mod.profiledAlloc"},
+	}
+	for _, tt := range tests {
+		if got := memProfileFunctionName(tt.name); got != tt.want {
+			t.Fatalf("memProfileFunctionName(%q) = %q, want %q", tt.name, got, tt.want)
+		}
+	}
+}
+
+func TestCompileRuntimeMemProfileInstrumentation(t *testing.T) {
+	_, m := mustCompileLLPkgFromSrc(t, `package foo
+import "runtime"
+
+var _ = runtime.MemProfileRate
+
+func sample() {
+}
+`)
+	fn := mustNamedFunction(t, m, "foo.sample")
+	fnIR := fn.String()
+	for _, want := range []string{"MemProfileEnter", "MemProfileExit"} {
+		if !strings.Contains(fnIR, want) {
+			t.Fatalf("compiled function missing %s instrumentation:\n%s", want, fnIR)
+		}
+	}
+	ir := m.String()
+	for _, want := range []string{"noinline", "disable-tail-calls"} {
+		if !strings.Contains(ir, want) {
+			t.Fatalf("compiled module missing %s attribute:\n%s", want, ir)
+		}
+	}
+}