From a683881eb7d5474e2e5389d6f491167e207aa1d1 Mon Sep 17 00:00:00 2001 From: Jake Bailey <5341706+jakebailey@users.noreply.github.com> Date: Fri, 12 Jun 2026 11:22:07 -0700 Subject: [PATCH] transform: handle returned pointer alloc aliases Follow returned pointer aliases when deciding whether runtime.alloc calls can be lowered to stack allocations. A returned parameter is not the same as nocapture: it still flows back to the caller and must be checked as an alias of the original allocation. Keep the analysis conservative for recursive returned-parameter chains and unknown operands. The existing golden test now shows that the non-escaping returned pointer cases no longer require heap allocation. --- transform/allocs.go | 74 +++++++++++++++++++++++++-- transform/testdata/allocs2.out.cover | 2 - transform/testdata/allocs2.out.reason | 4 +- 3 files changed, 71 insertions(+), 9 deletions(-) diff --git a/transform/allocs.go b/transform/allocs.go index 477afa0a3..04dd988e6 100644 --- a/transform/allocs.go +++ b/transform/allocs.go @@ -162,6 +162,21 @@ func FormatAllocCover(pos token.Position) string { // valueEscapesAt returns the instruction where the given value may escape and a // nil llvm.Value if it definitely doesn't. The value must be an instruction. func valueEscapesAt(value llvm.Value) llvm.Value { + return valueEscapesAtImpl(value, false, nil) +} + +func valueEscapesAtImpl(value llvm.Value, allowReturn bool, visiting map[llvm.Value]struct{}) llvm.Value { + if visiting == nil { + visiting = make(map[llvm.Value]struct{}) + } + if _, ok := visiting[value]; ok { + // Recursive call graph while following returned parameters. Treat this + // as escaping to keep the analysis conservative and bounded. + return value + } + visiting[value] = struct{}{} + defer delete(visiting, value) + uses := getUses(value) for _, use := range uses { if use.IsAInstruction().IsNil() { @@ -169,12 +184,12 @@ func valueEscapesAt(value llvm.Value) llvm.Value { } switch use.InstructionOpcode() { case llvm.GetElementPtr: - if at := valueEscapesAt(use); !at.IsNil() { + if at := valueEscapesAtImpl(use, allowReturn, visiting); !at.IsNil() { return at } case llvm.BitCast: // A bitcast escapes if the casted-to value escapes. - if at := valueEscapesAt(use); !at.IsNil() { + if at := valueEscapesAtImpl(use, allowReturn, visiting); !at.IsNil() { return at } case llvm.Load: @@ -186,12 +201,16 @@ func valueEscapesAt(value llvm.Value) llvm.Value { return use } case llvm.Call: - if !hasFlag(use, value, "nocapture") { - return use + if at := callValueEscapesAt(use, value, allowReturn, visiting); !at.IsNil() { + return at } case llvm.ICmp: // Comparing pointers don't let the pointer escape. // This is often a compiler-inserted nil check. + case llvm.Ret: + if !allowReturn || use.Operand(0) != value { + return use + } default: // Unknown instruction, might escape. return use @@ -202,6 +221,53 @@ func valueEscapesAt(value llvm.Value) llvm.Value { return llvm.Value{} } +// callValueEscapesAt returns whether value escapes through this call. It also +// handles calls that return value unchanged, as long as the called function does +// not otherwise capture the parameter and the returned alias does not escape. +func callValueEscapesAt(call, value llvm.Value, allowReturn bool, visiting map[llvm.Value]struct{}) llvm.Value { + called := call.CalledValue() + if called.IsAFunction().IsNil() { + return call + } + kindNoCapture := llvm.AttributeKindID("nocapture") + kindReturned := llvm.AttributeKindID("returned") + matched := false + returned := false + for i := 0; i < called.ParamsCount(); i++ { + if call.Operand(i) != value { + continue + } + matched = true + index := i + 1 // param attributes start at 1 + nocapture := !called.GetEnumAttributeAtIndex(index, kindNoCapture).IsNil() + returnedParam := !called.GetEnumAttributeAtIndex(index, kindReturned).IsNil() + if returnedParam { + returned = true + } + if nocapture { + continue + } + if !returnedParam || called.IsDeclaration() { + return call + } + if at := valueEscapesAtImpl(called.Param(i), true, visiting); !at.IsNil() { + return at + } + } + for i := called.ParamsCount(); i < call.OperandsCount(); i++ { + if call.Operand(i) == value { + return call + } + } + if !matched { + return llvm.Value{} + } + if returned { + return valueEscapesAtImpl(call, allowReturn, visiting) + } + return llvm.Value{} +} + func lineLengthAt(filename string, lineNumber int) int { f, err := os.Open(filename) if err != nil { diff --git a/transform/testdata/allocs2.out.cover b/transform/testdata/allocs2.out.cover index 2d42ec7c8..4a45ae0d7 100644 --- a/transform/testdata/allocs2.out.cover +++ b/transform/testdata/allocs2.out.cover @@ -1,4 +1,3 @@ -testdata/allocs2.go:12.1,12.9 1 0 testdata/allocs2.go:21.1,21.22 1 0 testdata/allocs2.go:24.1,24.43 1 0 testdata/allocs2.go:26.1,26.25 1 0 @@ -8,7 +7,6 @@ testdata/allocs2.go:44.1,44.23 1 0 testdata/allocs2.go:46.1,46.22 1 0 testdata/allocs2.go:49.1,49.9 1 0 testdata/allocs2.go:50.1,50.9 1 0 -testdata/allocs2.go:98.1,98.23 1 0 testdata/allocs2.go:107.1,107.21 1 0 testdata/allocs2.go:114.1,114.23 1 0 testdata/allocs2.go:128.1,128.23 1 0 diff --git a/transform/testdata/allocs2.out.reason b/transform/testdata/allocs2.out.reason index 82c2efd97..6bb45f623 100644 --- a/transform/testdata/allocs2.out.reason +++ b/transform/testdata/allocs2.out.reason @@ -1,4 +1,3 @@ -testdata/allocs2.go:12:2: object allocated on the heap: escapes at line 13 testdata/allocs2.go:21:12: object allocated on the heap: escapes at line 22 testdata/allocs2.go:24:15: object allocated on the heap: size is not constant testdata/allocs2.go:26:12: object allocated on the heap: object size 300 exceeds maximum stack allocation size 256 @@ -8,7 +7,6 @@ testdata/allocs2.go:44:22: object allocated on the heap: escapes at line 44 testdata/allocs2.go:46:13: object allocated on the heap: escapes at line 47 testdata/allocs2.go:49:2: object allocated on the heap: escapes at line 51 testdata/allocs2.go:50:2: object allocated on the heap: escapes at line 51 -testdata/allocs2.go:98:2: object allocated on the heap: escapes at line 100 testdata/allocs2.go:107:11: object allocated on the heap: escapes at line 108 -testdata/allocs2.go:114:2: object allocated on the heap: escapes at line 116 +testdata/allocs2.go:114:2: object allocated on the heap: escapes at line 117 testdata/allocs2.go:128:2: object allocated on the heap: escapes at line 130