Zenn
🔁

Go の iterator が難しかったので理解をまとめる

2024/06/23に公開
Go
tech

Go 1.22で experimental な機能として iterator が実装されました（range over functionと呼ばれることもあります）。Go Wikiの例を引用すると、以下のように定義されるslices.Backward関数を使って

package slices

func Backward[E any](s []E) func(func(int, E) bool) {
        return func(yield func(int, E) bool) {
                for i := len(s)-1; i >= 0; i-- {
                        if !yield(i, s[i]) {
                                return
                        }
                }
        }
}

次のようなfor-rangeループを書くことでスライスsを逆順に走査することができます（Go Playground）。

func main() {
	s := []string{"hello", "world"}
	for i, x := range slices.Backward(s) {
		fmt.Println(i, x)
	}
        // Output:
        // 1 world
        // 0 hello
}

混乱

さて、これを見たとき、自分はどうしてこのような挙動になるのか理解できませんでした。いくつか記事を読んだり、コンパイラを覗いたりすることで自分なりの理解を確立したので、それをまとめます。あくまで自分の理解を整理したものなので、他の人にとって分かりやすいかは不明ですが、誰かの役に立てば幸いです。

まずは、具体的にどこが混乱を招いていたのかを整理します（以下、iteratorを使用したfor-rangeループをiteratorループと呼ぶことにします）。

混乱1: forが2つ出てくる

上のコード例には、main関数とslices.Backward関数の双方でfor文が登場します。なぜ1つのループのために2つのfor文が現れるのかわかりませんでした。

混乱2: 謎のyield引数

slices.Backward(s)はyieldという引数をとる関数ですが、yieldに何が渡されてこの関数がどう使われるのかよくわからなかったです。

理解

これを腹落ちさせるには、iteratorループをコンパイラがどう処理するかについて理解する必要がありました。具体的には、以下の2点です。

  1. Iteratorループ全体が1つの関数呼び出しに変換されること。
  2. ループ内の処理が関数に変換されて、1.の関数呼び出しの引数に代入されること。

Iteratorループ全体が1つの関数呼び出しに変換されること

Iteratorループは、単なる関数呼び出しです。例えば、上で例示したslices.Backwardを使うコードは、コンパイラによって以下のように変換されます（yield引数については後述。とにかくただの関数呼び出しであるということを強調しています）。

slices.Backward(s)(yield)

変換後の処理にはfor-rangeの面影はありません。ただslices.Backward(s)を呼んでるだけです。しかし実際にはsを逆向きに走査するループが実現できています。これは運良くslices.Backwardの内部でループを行ってくれているためです。極端な話、一切ループしないようなiteratorも書けます。例えば下のコードを実行すると、for-rangeの中身によらず「ループなんてしないよ」と1回だけ出力されて終了します（Go Playground）。これは当然、iteratorループがNoLoopを呼び出しているだけだからです。

package main

import (
	"fmt"
)

func NoLoop[E any](s []E) func(func(int, E) bool) {
	return func(yield func(int, E) bool) {
		fmt.Println("ループなんてしないよ")
	}
}

func main() {
	s := []string{"hello", "world"}
	for i, x := range NoLoop(s) {
		fmt.Println(i, x)
	}
        // Output:
        // ループなんてしないよ
}

ループ内の処理が関数に変換されて、yield引数に代入されること

前項で説明したように、iteratorループはただの関数呼び出しです。ではその引数にはなにが渡されるのでしょうか。

slices.Backward(s)(yield)

これは、iteratorループ内の処理が関数に変換されて渡されます。冒頭で挙げた以下のコードを例にとると、

for i, x := range slices.Backward(s) {
        fmt.Println(i, x)
}

このループ内のfmt.Println(i, x)が以下のように関数に変換されslices.Backward(s)に渡されます。yield関数が(int i, string x)を引数にとっていることは、iteratorループでfor i, x :=と2つの変数を受け取っていることに対応します。普通のfor-rangeのように変数を0個または1個だけ受け取ることも可能です。

yield := func(int i, string x) bool {
        fmt.Println(i, x)
        return true
}
slices.Backward(s)(yield)

関数に変換される際、ループ内のbreakはreturn falseに、continueはreturn trueに変換されます[1]。また上の例のように制御構文がない場合は末尾でreturn trueされます。なので、例えば次のループは、

s := []int{100, 200, 300, 400, 500}
for i, x := range slices.Backward(s) {
        if x < 300 {
                break
        }
        if x > 400 {
                continue
        }
        fmt.Println(i, x)
}

以下のように翻訳されて引数に渡されます。

yield := func(i int, x int) bool {
        if x < 300 {
                return false
        }
        if x > 400 {
                return true
        }
        fmt.Println(i, x)
        return true
}
slices.Backward(s)(yield)
// Output:
// 3 400
// 2 300

slices.Backwardの実装を見ると、yield(i, s[i])の返り値がfalseの場合は運良く直ちにreturnしてくれています。これによって、iteratorループ内でbreakしたら、実際にループを行っているslices.Backwardのループも中断されます。

func Backward[E any](s []E) func(func(int, E) bool) {
    return func(yield func(int, E) bool) {
        for i := len(s)-1; i >= 0; i-- {
            if !yield(i, s[i]) {
                return
            }
        }
    }
}

最後に

ここまで見たように、iteratorループは実際には「ループ」とは限りません。for-rangeがただの関数呼び出しに翻訳されたり、ループ内の処理が関数に翻訳されたりするのは、Goの他構文と比較してややsurprisingに感じました。

逆に、iteratorを実装する側は、利用者から見て「ただのfor-range」に見えるように実装することを心掛けるべきと考えます。具体的には、slices.Backwardで見たようにyieldが値を返した場合（= iteratorループ内でcontinue, breakしたり、ループの末尾に到達した場合）関数を直ちにreturnすることや、状態をmutateせずfor-rangeに値を供給する責務に専念することなどがあるかと思います。

脚注

  1. Labeled breakなどを考慮するとそんなに単純ではないようですが、ここでは説明のため簡略化します。詳細は https://github.com/golang/go/blob/master/src/cmd/compile/internal/rangefunc/rewrite.go 冒頭のコメントに記述されています。 ↩︎

Discussion