cmd/compile: declare and assign of function literal escapes to heap

[qdongxu]

Go version

1.23.2

Output of go env in your module/workspace:

go env
GO111MODULE='on'
GOARCH='arm64'
GOBIN=''
GOCACHE='/Users/user/Library/Caches/go-build'
GOENV='/Users/user/Library/Application Support/go/env'
GOEXE=''
GOEXPERIMENT=''
GOFLAGS=''
GOHOSTARCH='arm64'
GOHOSTOS='darwin'
GOINSECURE=''
GOMODCACHE='/Users/user/go/pkg/mod'
GONOPROXY=''
GONOSUMDB=''
GOOS='darwin'
GOPATH='/Users/user/go'
GOPRIVATE=''
GOPROXY='https://goproxy.cn,direct'
GOROOT='/opt/homebrew/opt/go/libexec'
GOSUMDB='sum.golang.org'
GOTMPDIR=''
GOTOOLCHAIN='local'
GOTOOLDIR='/opt/homebrew/opt/go/libexec/pkg/tool/darwin_arm64'
GOVCS=''
GOVERSION='go1.23.2'
GODEBUG=''
GOTELEMETRY='local'
GOTELEMETRYDIR='/Users/user//Library/Application Support/go/telemetry'
GCCGO='gccgo'
GOARM64='v8.0'
AR='ar'
CC='cc'
CXX='c++'
CGO_ENABLED='1'
GOMOD='/Users/uesr/github/gorecycler/go.mod'
GOWORK=''
CGO_CFLAGS='-O2 -g'
CGO_CPPFLAGS=''
CGO_CXXFLAGS='-O2 -g'
CGO_FFLAGS='-O2 -g'
CGO_LDFLAGS='-O2 -g'
PKG_CONFIG='pkg-config'
GOGCCFLAGS='-fPIC -arch arm64 -pthread -fno-caret-diagnostics -Qunused-arguments -fmessage-length=0 -ffile-prefix-map=/var/folders/5b/483w2_yd7wn5x962q5hvrzkc0000gn/T/go-build271930357=/tmp/go-build -gno-record-gcc-switches -fno-common'

What did you do?

when declare an closure and assign in separate lines the local variable escape:

        var someFunc func(assignF func())
	someFunc = func(assignF func()) {
		assignF()
	}
	num := testing.AllocsPerRun(1, func() {
		i := 0 // i escaped
		someFunc(func() {
			x := i
			if x < 0 {
				// do nothing, just make x valid
			}
		})

	})

but when declare the closure in one statement, it does not escape:

       var someFunc = func(assignF func()) {
		assignF()
	}
	num := testing.AllocsPerRun(1, func() {
		i := 0 // i does not escaped
		someFunc(func() {
			x := i
			if x < 0 {
				// do nothing, just make x valid
			}
		})

	})

escape : https://go.dev/play/p/r36WtbRaJtf
not escape: https://go.dev/play/p/2u-ydl00_0i

What did you see happen?

The local variable i does not escape in one code block but escaped in another code block.

What did you expect to see?

The variable i should not escape in both code blocks.

[gabyhelp]

Related Issues and Documentation

• cmd/compile: escape analysis on closures #39511 (closed)
• Go escape analysis doesn't report variable escape with slice append #45148 (closed)
• cmd/compile: backing a local slice with an array in a loop causes the array to escape #24730 (closed)
• cmd/compile: argument escapes when it should not #18001 (closed)
• cmd/compile: unnecessarily moving variables into heap when using goroutine #33216 (closed)
• cmd/compile: Pointer escapes when stored in a non-escaping map #63343 (closed)
• cmd/compile: unexpected heap allocation passing pointer to non-escaping closure #14699 (closed)
• cmd/compile: Autotmp* variables causing unnecessary escape #58111 (closed)
• cmd/compile: generic function argument causes escape to heap #48849

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

[seankhliao]

If you look at the output of go build -gcflags '-m=2', it's the function literal that escapes, not the local variable.

[randall77]

I think this is probably the TODO here:

go/src/cmd/compile/internal/ir/expr.go

Line 930 in 989eed2

// TODO: handle initial declaration not including an assignment and

[qdongxu]

Probable the local variable is a trigger. There's no allocs reported after moving the i into the assignF inline function : https://go.dev/play/p/R6203vmDmlo

And the original observation is from some real project code. I did not reproduced in pay ground yet after simplifying the code. but the allocs not happened after the local variable moved or removed.

Senario 1:

func NewNode() recycle.Recycler[producer.Node] {
	root := recycle.Get[producer.Node, node]()
	// i := 3 // the benchmark reports a alloc
	root.Assign(func(bt producer.Node) {
                 i := 3 // the benchmark does not reports allocs after `i` moved here
		newNode(bt.(*node), i)
	})

	return root
}

Senario 2:
Benchmark reports 2 allocs:

func newNode(n *node, layer int) {
	if layer <= 0 {
		return
	}

	lRecycler := recycle.Get[producer.Node, node]()
	rRecycler := recycle.Get[producer.Node, node]()
	var lNode, rNode producer.Node // benchmark reports 2 allocs 
	lRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer
		n.left = bt
		lNode = bt
	})
	rRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer 
		n.right = bt
		rNode = bt
	})

	newNode(lNode.(*node), layer-1)
	newNode(rNode.(*node), layer-1)
}

benchmark reports no allocs after removing the local variables lNode, rNode:

func newNode(n *node, layer int) {
	if layer <= 0 {
		return
	}

	lRecycler := recycle.Get[producer.Node, node]()
	rRecycler := recycle.Get[producer.Node, node]()
	lRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer
		n.left = bt
		newNode(bt.(*node), layer-1)
	})
	rRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer
		n.right = bt
		newNode(bt.(*node), layer-1)
	})
}

the allocs increase accordingly after adding additional variables lNode2 and rNode2:

func newNode(n *node, layer int) {
	if layer <= 0 {
		return
	}

	lRecycler := recycle.Get[producer.Node, node]()
	rRecycler := recycle.Get[producer.Node, node]()
	var lNode, rNode producer.Node
	var lNode2, rNode2 producer.Node
	lRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer + 1
		n.left = bt
		lNode = bt
		lNode2 = bt
	})
	rRecycler.Assign(func(bt producer.Node) {
		bt.(*node).i = layer + 1
		n.right = bt
		rNode = bt
		rNode2 = bt
	})

	newNode(lNode.(*node), layer-1)
	newNode(rNode.(*node), layer-1)
	newNode(lNode2.(*node), layer-1)
	newNode(rNode2.(*node), layer-1)
}

[qdongxu]

I verified with go build -gcflags "-m=2" . > 1.log 2>&1 and confirmed that the i does not escape to heap. It's the closure literal that escape to heap. But the lNode, rNode lNode2, rNode2 do escape:

./consumer.go:25:6: newNode capturing by ref: lNode (addr=false assign=true width=16)
./consumer.go:25:6: lNode escapes to heap:
./consumer.go:25:6:   flow: {storage for func literal} = &lNode:
./consumer.go:25:6:     from lNode (captured by a closure) at ./consumer.go:30:3
./consumer.go:25:6:     from lNode (reference) at ./consumer.go:30:3
./consumer.go:26:6: newNode capturing by ref: lNode2 (addr=false assign=true width=16)
./consumer.go:26:6: lNode2 escapes to heap:
./consumer.go:26:6:   flow: {storage for func literal} = &lNode2:
./consumer.go:26:6:     from lNode2 (captured by a closure) at ./consumer.go:31:3
./consumer.go:26:6:     from lNode2 (reference) at ./consumer.go:31:3
./consumer.go:18:23: newNode capturing by value: layer (addr=false assign=false width=8)
./consumer.go:18:14: newNode capturing by value: n (addr=false assign=false width=8)
./consumer.go:25:13: newNode capturing by ref: rNode (addr=false assign=true width=16)
./consumer.go:25:13: rNode escapes to heap:
./consumer.go:25:13:   flow: {storage for func literal} = &rNode:
./consumer.go:25:13:     from rNode (captured by a closure) at ./consumer.go:36:3
./consumer.go:25:13:     from rNode (reference) at ./consumer.go:36:3
./consumer.go:26:14: newNode capturing by ref: rNode2 (addr=false assign=true width=16)
./consumer.go:26:14: rNode2 escapes to heap:
./consumer.go:26:14:   flow: {storage for func literal} = &rNode2:
./consumer.go:26:14:     from rNode2 (captured by a closure) at ./consumer.go:37:3
./consumer.go:26:14:     from rNode2 (reference) at ./consumer.go:37:3

[cherrymui]

cc @dr2chase for how we handle closures (with a declaration and then assignment, similar to the autotmp case).