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

@@ -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 {

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

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