【Go】rangefuncでループ終了後にfor文本体を呼び出すとエラーになる

TL; DR

var f func(string) bool

func capture(yield func(string) bool) {
	f = yield
}

func main() {
	for s := range capture {
		fmt.Println(s)
	}
	f("foo")
}

panic: runtime error: range function continued iteration after exit

はじめに

Go1.22のexperimentalの機能として、for文の range に関数が使用できるようになりました。
イテレータとして値を生成しながらループをする場合などに便利です。

func main() {
	// 関数が使える！
	for i := range Fact {
		if i > 1000000 {
			break
		}
		fmt.Println(i)
	}
}

1
2
6
24
120
720
5040
40320
362880

func Fact(yield func(int) bool) {
	fact := 1
	i := 1
	for {
		// yieldを呼ぶたびに、引数がイテレーション変数に渡されfor文の本体が実行される
		// breakやreturn等でfor文を抜けた場合、戻り値にfalseが返る
		if !yield(fact) {
			return
		}
		i++
		fact *= i
	}
}

yieldを呼べる条件

ところで、for文本体を yield で呼び出せるのであれば、yield を保持しておけばループを抜けた後でもfor文の本体を実行できるのではないでしょうか？
...この目論見は冒頭の通り失敗しました。

再掲

var f func(string) bool

func capture(yield func(string) bool) {
	// グローバル変数に代入
	f = yield
}

func main() {
	for s := range capture {
		fmt.Println(s)
	}
	f("foo") // panic: runtime error: range function continued iteration after exit
}

エラーメッセージの通り、ループを抜けた後に yield を呼ぼうとすると実行時にpanicします。

To permit checking that an iterator is well-behaved -- that is, that it does not call the loop body again after it has returned false or after the entire loop has exited (it might retain a copy of the body function, or pass it to another goroutine) -- each generated loop has its own #exitK flag that is checked before each iteration, and set both at any early exit and after the iteration completes.

イテレータが正しくふるまう―言い換えると、イテレータが false を返したりループ全体が終了したりした後にはループボディが呼び出されない（イテレータはボディ関数のコピーを保持していたり、別のゴルーチンに渡してしまうかもしれない）―ことをチェックするため、生成された各ループは個別の #exitK フラグを持ち、各イテレーションの前にチェックされ、早期終了やイテレーション完了後の両方で設定されます。

内部的に以下のように置き換えられるという記載もありました。

before

for x := range f {
	...
	if ... { break }
	...
}

after

{
	var #exit1 bool
	f(func(x T1) bool {
		if #exit1 { runtime.panicrangeexit() }
		...
		if ... { #exit1 = true ; return false }
		...
		return true
	})
	#exit1 = true
}

というわけで、黒魔術には応用できないようです。悪いことができないようになっていて安心ですね