Avoid heap escapes with iter.Seq

I have a []byte with backing array allocated on the stack I would like to use an iterator to go over it. However, I cannot figure how to avoid heap escape. go build -gcflags="-m -m" of the following snippet

package main

import "iter"

func everyByte(data []byte) iter.Seq[[]byte] {
	return func(yield func([]byte) bool) {
		for i := range data {
			if !yield(data[i : i+1]) {
				break
			}
		}
	}
}

func main() {
	var x int
	dataStore := [...]byte{1, 2, 3, 4, 5, 6, 7}
	data := dataStore[:]
	for b := range everyByte(data) {
		x += int(b[0])
	}
	println(x)
}

shows that

./main.go:17:2: dataStore escapes to heap in main:
./main.go:17:2:   flow: data ← &dataStore:
./main.go:17:2:     from dataStore (address-of) at ./main.go:18:19
./main.go:17:2:     from dataStore[:] (slice) at ./main.go:18:19
./main.go:17:2:     from data := dataStore[:] (assign) at ./main.go:18:7
./main.go:17:2:   flow: data ← data:
./main.go:17:2:     from data := data (assign-pair) at ./main.go:19:26
./main.go:17:2:   flow: {heap} ← data:
./main.go:17:2:     from data[i:i + 1] (slice) at ./main.go:8:18
./main.go:17:2:     from yield(data[i:i + 1]) (call parameter) at ./main.go:8:13
./main.go:6:14: yield does not escape
./main.go:17:2: moved to heap: dataStore

Is it possible to rewrite it to keep dataStore on the stack? Iterator is so much nicer than maintaining an offset and having to call “NextBlock”-like function.

Hello. This value escapes to the heap, because you return a slice from iterator. Compiler cannot guarantee its lifetime after returning and thus, allocates it on the heap. If you switch iteration over every byte, not slice, then it’s stays on the stack.

P.S. calling an iterator everyByte and then returning a slice is very misleading.

1 Like

Thanks @lemarkar . This is just to set up an example. In the real code I parse bytes into “message” blocks.

You gave me an idea though. Instead of returning a proper slice, I can return lo:hi and turn it into a slice when needed.

package main

import "iter"

func everyMessage(data []byte) iter.Seq2[int, int] {
	return func(yield func(int, int) bool) {
		for i := range data {
			if !yield(i, i+1) {
				break
			}
		}
	}
}

func main() {
	dataStore := [...]byte{1, 2, 3, 4, 5, 6, 7}
	data := dataStore[:]
	for lo, hi := range everyMessage(data) {
		msg := data[lo:hi]
		println(msg)
	}
}

Could this also be fixed by persuading the compiler to inline this function ? If the function is inlined the lifetime of the slice should be obvious and no escape necessary ?