transform: various escape analysis improvements

transform/allocs.go

@@ -162,44 +162,141 @@ 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).escapeAt
+}
+
+type escapeResult struct {
+	escapeAt llvm.Value
+
+	// returned is separate from escapeAt because values can flow to a return
+	// through aggregate operations. LLVM can mark a scalar returned parameter,
+	// but a slice data pointer returned inside {ptr, len, cap} is only visible
+	// after walking insertvalue/ret uses in the callee.
+	returned bool
+}
+
+func (r *escapeResult) merge(other escapeResult) bool {
+	if !other.escapeAt.IsNil() {
+		r.escapeAt = other.escapeAt
+		return false
+	}
+	r.returned = r.returned || other.returned
+	return true
+}
+
+func valueEscapesAtImpl(value llvm.Value, allowReturn bool, visiting map[llvm.Value]struct{}) escapeResult {
+	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 escapeResult{escapeAt: value}
+	}
+	visiting[value] = struct{}{}
+	defer delete(visiting, value)
+
+	var result escapeResult
 	uses := getUses(value)
 	for _, use := range uses {
 		if use.IsAInstruction().IsNil() {
 			panic("expected instruction use")
 		}
 		switch use.InstructionOpcode() {
 		case llvm.GetElementPtr:
-			if at := valueEscapesAt(use); !at.IsNil() {
-				return at
+			if !result.merge(valueEscapesAtImpl(use, allowReturn, visiting)) {
+				return result
 			}
 		case llvm.BitCast:
 			// A bitcast escapes if the casted-to value escapes.
-			if at := valueEscapesAt(use); !at.IsNil() {
-				return at
+			if !result.merge(valueEscapesAtImpl(use, allowReturn, visiting)) {
+				return result
+			}
+		case llvm.InsertValue:
+			if !result.merge(valueEscapesAtImpl(use, allowReturn, visiting)) {
+				return result
+			}
+		case llvm.ExtractValue:
+			if use.Type().TypeKind() == llvm.PointerTypeKind {
+				if !result.merge(valueEscapesAtImpl(use, allowReturn, visiting)) {
+					return result
+				}
 			}
 		case llvm.Load:
 			// Load does not escape.
 		case llvm.Store:
 			// Store only escapes when the value is stored to, not when the
 			// value is stored into another value.
 			if use.Operand(0) == value {
-				return use
+				return escapeResult{escapeAt: use}
 			}
 		case llvm.Call:
-			if !hasFlag(use, value, "nocapture") {
-				return use
+			if !result.merge(callValueEscapesAt(use, value, allowReturn, visiting)) {
+				return result
 			}
 		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 escapeResult{escapeAt: use}
+			}
+			result.returned = true
 		default:
 			// Unknown instruction, might escape.
-			return use
+			return escapeResult{escapeAt: use}
 		}
 	}
 
 	// Checked all uses, and none let the pointer value escape.
-	return llvm.Value{}
+	return result
+}
+
+// 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{}) escapeResult {
+	called := call.CalledValue()
+	if called.IsAFunction().IsNil() {
+		return escapeResult{escapeAt: call}
+	}
+	kindNoCapture := llvm.AttributeKindID("nocapture")
+	kindReturned := llvm.AttributeKindID("returned")
+	matched := false
+	var result escapeResult
+	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 {
+			result.returned = true
+		}
+		if nocapture {
+			continue
+		}
+		if called.IsDeclaration() {
+			return escapeResult{escapeAt: call}
+		}
+		if !result.merge(valueEscapesAtImpl(called.Param(i), true, visiting)) {
+			return result
+		}
+	}
+	for i := called.ParamsCount(); i < call.OperandsCount(); i++ {
+		if call.Operand(i) == value {
+			return escapeResult{escapeAt: call}
+		}
+	}
+	if !matched {
+		return escapeResult{}
+	}
+	if result.returned {
+		return valueEscapesAtImpl(call, allowReturn, visiting)
+	}
+	return escapeResult{}
 }
 
 func lineLengthAt(filename string, lineNumber int) int {

transform/testdata/allocs2.go

@@ -9,7 +9,6 @@ func main() {
 	n1 := 5
 	derefInt(&n1)
 
-	// This should eventually be modified to not escape.
 	n2 := 6
 	returnIntPtr(&n2)
 
@@ -19,7 +18,6 @@ func main() {
 	s2 := [3]int{}
 	readIntSlice(s2[:])
 
-	// This should also be modified to not escape.
 	s3 := make([]int, 3)
 	returnIntSlice(s3)
 
@@ -78,6 +76,61 @@ func main() {
 	keepAliveNoEscape(unsafe.Pointer(&dmaBuf2[0]))
 }
 
+type vector3 [3]float32
+
+func scaleVector3(vec *vector3, f float32) *vector3 {
+	vec[0] *= f
+	vec[1] *= f
+	vec[2] *= f
+	return vec
+}
+
+func crossVector3(a, b *vector3) vector3 {
+	return vector3{
+		a[1]*b[2] - a[2]*b[1],
+		a[2]*b[0] - a[0]*b[2],
+		a[0]*b[1] - a[1]*b[0],
+	}
+}
+
+func nonEscapingReturnedPointer() vector3 {
+	a := vector3{1, 2, 3}
+	b := vector3{4, 5, 6}
+
+	c := scaleVector3(&b, 0.5)
+	return crossVector3(&a, c)
+}
+
+var escapedSlice []int
+
+func escapingReturnedSlice() {
+	s := make([]int, 3)
+	escapedSlice = returnIntSlice(s)
+}
+
+var escapedVector3 *vector3
+
+func escapingReturnedPointer() {
+	b := vector3{4, 5, 6}
+
+	c := scaleVector3(&b, 0.5)
+	escapedVector3 = c
+}
+
+func recursiveScaleVector3(vec *vector3, n int) *vector3 {
+	if n == 0 {
+		return vec
+	}
+	return recursiveScaleVector3(vec, n-1)
+}
+
+func recursiveReturnedPointer() vector3 {
+	b := vector3{4, 5, 6}
+
+	c := recursiveScaleVector3(&b, 1)
+	return *c
+}
+
 func derefInt(x *int) int {
 	return *x
 }

transform/testdata/allocs2.out.cover

@@ -1,10 +1,11 @@
-testdata/allocs2.go:13.1,13.9 1 0
-testdata/allocs2.go:23.1,23.22 1 0
-testdata/allocs2.go:26.1,26.43 1 0
-testdata/allocs2.go:28.1,28.25 1 0
-testdata/allocs2.go:31.1,31.22 1 0
-testdata/allocs2.go:38.1,38.23 1 0
-testdata/allocs2.go:46.1,46.23 1 0
-testdata/allocs2.go:48.1,48.22 1 0
-testdata/allocs2.go:51.1,51.9 1 0
-testdata/allocs2.go:52.1,52.9 1 0
+testdata/allocs2.go:24.1,24.43 1 0
+testdata/allocs2.go:26.1,26.25 1 0
+testdata/allocs2.go:29.1,29.22 1 0
+testdata/allocs2.go:36.1,36.23 1 0
+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: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,10 +1,11 @@
-testdata/allocs2.go:13:2: object allocated on the heap: escapes at line 14
-testdata/allocs2.go:23:12: object allocated on the heap: escapes at line 24
-testdata/allocs2.go:26:15: object allocated on the heap: size is not constant
-testdata/allocs2.go:28:12: object allocated on the heap: object size 300 exceeds maximum stack allocation size 256
-testdata/allocs2.go:31:12: object allocated on the heap: escapes at line 32
-testdata/allocs2.go:38:21: object allocated on the heap: escapes at line 39
-testdata/allocs2.go:46:22: object allocated on the heap: escapes at line 46
-testdata/allocs2.go:48:13: object allocated on the heap: escapes at line 49
-testdata/allocs2.go:51:2: object allocated on the heap: escapes at line 53
-testdata/allocs2.go:52:2: object allocated on the heap: escapes at line 53
+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
+testdata/allocs2.go:29:12: object allocated on the heap: escapes at line 30
+testdata/allocs2.go:36:21: object allocated on the heap: escapes at line 37
+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:107:11: object allocated on the heap: escapes at line 108
+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 unknown line