transform: track escapes through nested aggregate results

testdata/calls.go

@@ -1,5 +1,7 @@
 package main
 
+import "runtime/volatile"
+
 type Thing struct {
 	name string
 }
@@ -87,6 +89,180 @@ func main() {
 	if testDeferElse(false) != 0 {
 		println("else defer returned wrong value")
 	}
+
+	// See https://github.com/tinygo-org/tinygo/pull/5776.
+	err := multiReturnEscape()
+	clobberStack()
+	println("multiple return values:", err.(*multiReturnErr).x[1])
+	err = singleArrayEscape()
+	clobberStack()
+	println("single array return:", err.(*multiReturnErr).x[1])
+	err = forwardedReturnEscape()
+	clobberStack()
+	println("forwarded return:", err.(*multiReturnErr).x[1])
+	s := sliceReturnEscape()
+	clobberStack()
+	println("slice return:", s[1])
+	globalReturnEscape()
+	clobberStack()
+	println("global escape:", globalEscaped.(*multiReturnErr).x[1])
+	fn := closureReturnEscape()
+	clobberStack()
+	println("closure escape:", fn())
+	var phiCond uint32
+	err = phiReturnEscape(volatile.LoadUint32(&phiCond) != 0)
+	clobberStack()
+	println("phi escape:", err.(*multiReturnErr).x[1])
+	volatile.StoreUint32(&phiCond, 1)
+	println("phi nil:", phiReturnEscape(volatile.LoadUint32(&phiCond) != 0) == nil)
+	g := interfaceReturnEscape()
+	clobberStack()
+	println("interface return:", g.get())
+	err = deepForwardEscape()
+	clobberStack()
+	println("deep forward:", err.(*multiReturnErr).x[1])
+}
+
+type multiReturnErr struct{ x [4]int }
+
+func (e *multiReturnErr) Error() string { return "multiReturnErr" }
+
+//go:noinline
+func wrapMultiReturn(e *multiReturnErr) (int, error) { return 1, e }
+
+//go:noinline
+func multiReturnEscape() error {
+	e := &multiReturnErr{}
+	e.x[1] = 42
+	_, err := wrapMultiReturn(e)
+	return err
+}
+
+//go:noinline
+func wrapSingleArrayReturn(e *multiReturnErr) [1]error { return [1]error{e} }
+
+//go:noinline
+func singleArrayEscape() error {
+	e := &multiReturnErr{}
+	e.x[1] = 43
+	a := wrapSingleArrayReturn(e)
+	return a[0]
+}
+
+//go:noinline
+func forwardMultiReturn(e *multiReturnErr) error {
+	_, err := wrapMultiReturn(e)
+	return err
+}
+
+//go:noinline
+func forwardedReturnEscape() error {
+	e := &multiReturnErr{}
+	e.x[1] = 44
+	return forwardMultiReturn(e)
+}
+
+//go:noinline
+func wrapSliceReturn(b []int) (int, []int) { return 1, b }
+
+// The allocation reaches the caller inside the {ptr, len, cap} slice aggregate
+// of a multiple return value.
+//
+//go:noinline
+func sliceReturnEscape() []int {
+	b := make([]int, 4)
+	b[1] = 45
+	_, s := wrapSliceReturn(b)
+	return s
+}
+
+var globalEscaped error
+
+// The allocation escapes to a package-level variable rather than through a
+// return.
+//
+//go:noinline
+func globalReturnEscape() {
+	e := &multiReturnErr{}
+	e.x[1] = 46
+	_, err := wrapMultiReturn(e)
+	globalEscaped = err
+}
+
+// The allocation escapes by being captured in a closure.
+//
+//go:noinline
+func closureReturnEscape() func() int {
+	e := &multiReturnErr{}
+	e.x[1] = 47
+	_, err := wrapMultiReturn(e)
+	return func() int { return err.(*multiReturnErr).x[1] }
+}
+
+//go:noinline
+func phiReturnEscape(cond bool) error {
+	e := &multiReturnErr{}
+	e.x[1] = 48
+	_, err := wrapMultiReturn(e)
+	if cond {
+		err = nil
+	}
+	return err
+}
+
+type multiReturnGetter interface{ get() int }
+
+func (e *multiReturnErr) get() int { return e.x[1] }
+
+//go:noinline
+func wrapInterfaceReturn(e *multiReturnErr) (int, multiReturnGetter) { return 1, e }
+
+// The allocation reaches the caller inside a non-error interface value and is
+// then called through.
+//
+//go:noinline
+func interfaceReturnEscape() multiReturnGetter {
+	e := &multiReturnErr{}
+	e.x[1] = 49
+	_, g := wrapInterfaceReturn(e)
+	return g
+}
+
+//go:noinline
+func deepForward1(e *multiReturnErr) (error, int) {
+	_, err := wrapMultiReturn(e)
+	return err, 2
+}
+
+//go:noinline
+func deepForward2(e *multiReturnErr) [1]error {
+	err, _ := deepForward1(e)
+	return [1]error{err}
+}
+
+//go:noinline
+func deepForward3(e *multiReturnErr) error {
+	a := deepForward2(e)
+	return a[0]
+}
+
+// The allocation is repacked into a different aggregate shape at each of four
+// call levels before reaching the caller.
+//
+//go:noinline
+func deepForwardEscape() error {
+	e := &multiReturnErr{}
+	e.x[1] = 50
+	return deepForward3(e)
+}
+
+//go:noinline
+func clobberStack() {
+	var a [32]uint32
+	for i := range a {
+		// Keep the stack writes. See https://llvm.org/docs/LangRef.html#volatile-memory-accesses.
+		volatile.StoreUint32(&a[i], uint32(1000+i))
+	}
 }
 
 func runFunc(f func(int), arg int) {

testdata/calls.txt

@@ -16,3 +16,13 @@ thing inside closure: foo
 inside fp closure: foo 3
 Thing.Print:  arg: functional args 1
 Thing.Print: named thing arg: functional args 2
+multiple return values: 42
+single array return: 43
+forwarded return: 44
+slice return: 45
+global escape: 46
+closure escape: 47
+phi escape: 48
+phi nil: true
+interface return: 49
+deep forward: 50

transform/allocs.go

@@ -218,7 +218,7 @@ func valueEscapesAtImpl(value llvm.Value, allowReturn bool, visiting map[llvm.Va
 				return result
 			}
 		case llvm.ExtractValue:
-			if use.Type().TypeKind() == llvm.PointerTypeKind {
+			if typeHasPointers(use.Type()) {
 				if !result.merge(valueEscapesAtImpl(use, allowReturn, visiting)) {
 					return result
 				}

transform/allocs_test.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/tinygo-org/tinygo/compiler/llvmutil"
 	"github.com/tinygo-org/tinygo/transform"
 	"tinygo.org/x/go-llvm"
 )
@@ -20,6 +21,144 @@ func TestAllocs(t *testing.T) {
 	})
 }
 
+func TestAllocsAggregateEdges(t *testing.T) {
+	t.Parallel()
+
+	const path = "testdata/allocs-aggregate.ll"
+	ctx := llvm.NewContext()
+	defer ctx.Dispose()
+	ensureTestCacheFreshness(t, path)
+	buf, err := llvm.NewMemoryBufferFromFile(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	mod, err := ctx.ParseIR(buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer mod.Dispose()
+	if err := llvm.VerifyModule(mod, llvm.ReturnStatusAction); err != nil {
+		t.Fatal(err)
+	}
+
+	transform.OptimizeAllocs(mod, nil, 256, nil)
+	if err := llvm.VerifyModule(mod, llvm.ReturnStatusAction); err != nil {
+		t.Fatal(err)
+	}
+
+	for name, wantHeap := range map[string]bool{
+		"nestedStructReturn":  true,
+		"nestedArrayReturn":   true,
+		"arrayOfStructReturn": true,
+		"structOfArrayReturn": true,
+		"singleArrayReturn":   true,
+		"aggregateStore":      true,
+		"pointerStore":        true,
+		"aggregateCall":       true,
+		"unknownCall":         true,
+		"indirectCall":        true,
+		"repackReturn":        true,
+		"forwardReturn":       true,
+		"forwardLoad":         false,
+		"gepReturn":           true,
+		"gepLoad":             false,
+		"aggregatePhi":        true,
+		"aggregateSelect":     true,
+		"recursiveReturn":     true,
+		"duplicateArguments":  true,
+		"nonEscapingScalar":   false,
+		"nonEscapingLoad":     false,
+		"nonEscapingNilCheck": false,
+		"nonEscapingDiscard":  false,
+		"nonEscapingCall":     false,
+		"pointerFreeStruct":   false,
+		"pointerFreeArray":    false,
+		"emptyPointerArray":   false,
+	} {
+		t.Run(name, func(t *testing.T) {
+			fn := mod.NamedFunction(name)
+			if fn.IsNil() {
+				t.Fatal("function not found")
+			}
+			var heap, stack int
+			for bb := fn.FirstBasicBlock(); !bb.IsNil(); bb = llvm.NextBasicBlock(bb) {
+				for inst := bb.FirstInstruction(); !inst.IsNil(); inst = llvm.NextInstruction(inst) {
+					if !inst.IsACallInst().IsNil() && inst.CalledValue() == mod.NamedFunction("runtime.alloc") {
+						heap++
+					}
+					if !inst.IsAAllocaInst().IsNil() {
+						stack++
+					}
+				}
+			}
+			want := 0
+			if wantHeap {
+				want = 1
+			}
+			if heap != want || stack != 1-want {
+				t.Errorf("got %d heap and %d stack allocations, want %d and %d:\n%s", heap, stack, want, 1-want, fn.String())
+			}
+		})
+	}
+}
+
+func TestAllocsRuntimePhi(t *testing.T) {
+	t.Parallel()
+
+	mod := compileGoFileForTesting(t, "../testdata/calls.go")
+	defer mod.Context().Dispose()
+	defer mod.Dispose()
+	po := llvm.NewPassBuilderOptions()
+	defer po.Dispose()
+	passes := "globalopt,ipsccp,instcombine,adce,function-attrs"
+	if llvmutil.Version() >= 18 {
+		passes = "globalopt,ipsccp,instcombine<no-verify-fixpoint>,adce,function-attrs"
+	}
+	if err := mod.RunPasses(passes, llvm.TargetMachine{}, po); err != nil {
+		t.Fatal(err)
+	}
+	fn := mod.NamedFunction("main.phiReturnEscape")
+	if fn.IsNil() {
+		t.Fatal("phiReturnEscape not found")
+	}
+	var returnedPhi llvm.Value
+	for bb := fn.FirstBasicBlock(); !bb.IsNil(); bb = llvm.NextBasicBlock(bb) {
+		for inst := bb.FirstInstruction(); !inst.IsNil(); inst = llvm.NextInstruction(inst) {
+			if !inst.IsAReturnInst().IsNil() && !inst.Operand(0).IsAPHINode().IsNil() {
+				returnedPhi = inst.Operand(0)
+			}
+		}
+	}
+	if returnedPhi.IsNil() {
+		t.Fatalf("runtime regression does not return a phi before escape analysis:\n%s", fn.String())
+	}
+	var extracted bool
+	for i := 0; i < returnedPhi.IncomingCount(); i++ {
+		if !returnedPhi.IncomingValue(i).IsAExtractValueInst().IsNil() {
+			extracted = true
+		}
+	}
+	if !extracted {
+		t.Fatalf("returned phi does not merge an extracted aggregate:\n%s", fn.String())
+	}
+
+	transform.OptimizeAllocs(mod, nil, 256, nil)
+	if err := llvm.VerifyModule(mod, llvm.ReturnStatusAction); err != nil {
+		t.Fatal(err)
+	}
+	var allocations int
+	for bb := fn.FirstBasicBlock(); !bb.IsNil(); bb = llvm.NextBasicBlock(bb) {
+		for inst := bb.FirstInstruction(); !inst.IsNil(); inst = llvm.NextInstruction(inst) {
+			if !inst.IsACallInst().IsNil() && inst.CalledValue() == mod.NamedFunction("runtime.alloc") {
+				allocations++
+			}
+		}
+	}
+	if allocations != 1 {
+		t.Fatalf("got %d heap allocations, want 1:\n%s", allocations, fn.String())
+	}
+}
+
 // Test with a Go file as input (for more accurate tests).
 func TestAllocs2(t *testing.T) {
 	t.Parallel()