Zenn
🥷

Go1.23のイテレータを完全理解する✌️

2024/06/20に公開
Go
イテレータ
tech

イテレータについて完全理解するぞ！！！！

皆さん、Go1.23でイテレータが実装されたことはご存知でしょうか？
Go1.18でジェネリクスが導入されて書き方の幅が広がったように、今回のイテレータもGoの書き方の幅を広げる機能になるのではと思っています！

このブログでは、イテレータ実装にあたって導入された型/関数とその使い方を深ぼっていきたいと思います！

このブログの中では「range over func」のことを「イテレータ」と呼ぶことにします。また、このブログ内で使用するコードはGo1.23.2で書かれています。手元でインストールして手を動かしながら、学んでいってもらえれば幸いです！

❯ go version
go version go1.23.2 darwin/amd64

さっそくイテレータを深ぼる

issueの内容をおさらい

おさらいとして、イテレータがどういう仕様なのかを見ていきます。
仕様はこちらのissueに一部まとまっているので抜粋します。最終的な仕様はThe Go Programming Language Specificationにまとめられていますが、issueの方がわかりやすいのでそちらを一部抜粋して説明します。

If f is a function type of the form func(yield func(T1, T2)bool) bool, then for x, y := range f { ... } is similar to f(func(x T1, y T2) bool { ... }), where the loop body has been moved into the function literal, which is passed to f as yield. The boolean result from yield indicates to f whether to keep iterating. The boolean result from f itself is ignored in this usage but present to allow easier composition of iterators.

(意訳)
関数 f が func(yield func(T1, T2) bool) bool の形式なら、for x, y := range f { ... } は f(func(x T1, y T2) bool { ... } に似ている。ここで、ループ本体は関数リテラルに移動し、yield として f に渡される。

yield からの真偽値の結果は f に反復を続けるかどうかを示す。この使用法では f それ自身からの真偽値の結果は無視されているが、イテレータの合成を容易にするために存在する。

I say "similar to" and not "completely equivalent to" because all the control flow statements in the loop body continue to have their original meaning. In particular, break, continue, defer, goto, and return all do exactly what they would do in range over a non-function.

(意訳)
私は「似ている」とは言ったが、「完全に等しい」とは言っていない。なぜなら、ループ本体内のすべての制御フロー文は、本来の意味を持ち続けるからだ。特に、break、continue、defer、goto、return はすべて、非関数の範囲内で行うのとまったく同じことを行う。

上記のissueでは以下の3つの形式についてイテレータが実装予定とあります。

  • func(func()bool) bool
  • func(func(V)bool) bool
  • func(func(K, V)bool) bool

しかし、Go1.23のリリースノートをみてみると、以下の3つの形式についてイテレータが可能となっています！

  • func(func()bool)
  • func(func(V)bool)
  • func(func(K, V)bool)

なので、このブログではリリースノートに沿った形での実装例を載せています！

コードを書いてみる

さて、実際にissueに書いてある内容を検証してみます！

issueの内容をまとめると、イテレータはざっくり以下のような性質を持つようです。

  1. ループ本体は関数リテラルに移動し、yield として f に渡される
  2. yield からの真偽値の結果は f に反復を続けるかどうかを示す
  3. yield からの真偽値の結果はイテレータの合成を容易にするために存在する
  4. ループ本体内のすべての制御フロー文 (break、continue、defer、goto、return) は本来の意味を持ち続ける

上記の性質を1つ1つ試して理解を深めてみたいと思います！

1. 「ループ本体は関数リテラルに移動し、yield として f に渡される」性質をみてみる

まずは簡単な例として、func(yield func() bool) の形式のイテレータで実験してみたいと思います。

func main() {
	for range f {
		fmt.Println("for range starts")
	}
}

func f(yield func() bool) {}

上記の書き方と「ループ本体は関数リテラルに移動し、yield として f に渡される」性質を考えると、f の関数の中で yield を呼ばないとループ本文が実行されない、と予想できます。 実際に動かしてみます。

❯ go run .
# 何も表示されない

予想通り何も表示されませんでした。

f の中で yield を呼んでみます。コードは以下のようになります。

 func f(yield func() bool) {
+	yield()
 }

実際に動かしてみると "for range starts" が表示されました 🎉

❯ go run .
for range starts

2. 「yield からの真偽値の結果は f に反復を続けるかどうかを示す」性質をみてみる

次は func(yield func(V) bool) の形式で見てみたいと思います。

まずは単純な例として値を受け取って出力するというコードを書いてみます。

func main() {
	for v := range f {
		fmt.Println(v)
	}
}

func f(yield func(int) bool) {
	for _, num := range []int{1, 2, 3, 4, 5} {
		yield(num)
	}
}

ちゃんと値が出力されますね。

❯ go run .
1
2
3
4
5

yield からの真偽値の結果を無視して、f の反復を停止させるために v が3の時に return するようにしてみます

 func main() {
 	for v := range f {
 		fmt.Println(v)
+		if v == 3 {
+			return
+		}
 	}
 }

実際に動かしてみるとpanicしてしまいました！

❯ go run .
1
2
3
panic: runtime error: range function continued iteration after function for loop body returned false

goroutine 1 [running]:
main.main-range1(0x3?)
        /Users/xxxx/repos/output-docs/golang/iterator/main.go:6 +0xc9
main.f(...)
        /Users/xxxx/repos/output-docs/golang/iterator/main.go:16
main.main()
        /Users/xxxx/repos/output-docs/golang/iterator/main.go:6 +0x8e
exit status 2

コンパイルエラーにはならずにランタイムでのパニック (panic: runtime error: range function continued iteration after exit ) が起こるようです。

「yield からの真偽値の結果は f に反復を続けるかどうかを示す」性質を加味すると、f の反復を停止させるためには yield の返り値を適切にハンドリングしてあげる必要があるようですね！

f を以下のように書き換えてみます。

 func f(yield func(int) bool) {
 	for _, num := range []int{1, 2, 3, 4, 5} {
+		if !yield(num) {
+			return
+		}
 	}
 }

実際に動かしてみると、確かに v = 3 の時で処理が停止しています🎉

❯ go run .
1
2
3

先ほどの例では if !yield(num) { ... } と書いてみましたが、if yield(num) { ... } でどういう挙動になるかも試すと理解が深まるかもしれません。

3. 「yield からの真偽値の結果はイテレータの合成を容易にするために存在する」性質を見てみる

「イテレータの合成」をみるために、「イテレータを受け取って、イテレータを返す」関数 g を作ってみます。
関数 g は、引数のイテレータから受け取ったintをstringに変換するイテレータを返す関数とします。

シグネチャとしては以下のようになると思います

func g(f func(yield func(int) bool)) func(yield func(string) bool) {
    // 実装
}

そして実装は以下のようにしてみます。

func g(f func(yield func(int) bool)) func(yield func(string) bool) {
	maps := map[int]string{
		1: "one", 2: "two", 3: "three",
		4: "four", 5: "five",
	}
	return func(yield func(string) bool) {
		for v := range f {
			str, ok := maps[v]
			if ok && !yield(str) {
				return
			}
		}
	}
}

先ほどの f 関数と組み合わせて、最終的なコードは以下のようになります。

func main() {
	for v := range g(f) {
		fmt.Println(v)
	}
}

func f(yield func(int) bool) {
	for _, num := range []int{1, 2, 3, 4, 5} {
		if !yield(num) {
			return
		}
	}
}

func g(f func(yield func(int) bool)) func(yield func(string) bool) {
	maps := map[int]string{
		1: "one", 2: "two", 3: "three",
		4: "four", 5: "five",
	}
	return func(yield func(string) bool) {
		for v := range f {
			str, ok := maps[v]
			if ok && !yield(str) {
				return
			}
		}
	}
}

実行してみます

❯ go run .
one
two
three
four
five

ちゃんとintがstringに変換されていますね！

前置きが長くなりましたが、ここで g の処理を途中でストップさせてみます。

 	return func(yield func(string) bool) {
 		for v := range f {
 			str, ok := maps[v]
+			if len(str) > 3 {
+				break
+			}
 			if ok && !yield(str) {
 				return
 			}
 		}
 	}

実行してみます

❯ go run .
one
two

当たり前といえば当たり前の結果なのですが、これが「yield からの真偽値の結果はイテレータの合成を容易にするために存在する」性質だと筆者は理解しています。

つまり、呼び出し側でどう処理されているかは不明だがyield の適切にハンドリングしていれば、イテレータ同士の連携が容易にできるということなのではないでしょうか

4. 「ループ本体内のすべての制御フロー文 (break、continue、defer、goto、return) は本来の意味を持ち続ける」性質を見てみる

先ほどのコードで既に見ているように、for loopの中の break や return の挙動は、イテレータを使う使わないに限らずそのままのようです

イテレータの挙動まとめ 🏃

今まで見てきたイテレータの挙動を軽くまとめてみます。

  1. ループ本文は yield としてイテレータに渡される
    • むしろ、yield が呼び出された時点で、処理がループ本文に戻ると考えた方が良いかもしれません
  2. yield から返り値である真偽値はイテレータを停止するかどうかを制御する
    • 真偽値を適切にハンドリングしないとランタイム時にパニックになる可能性があるので、個人的にはこの性質はとても重要なのではないかと思います
  3. yield からの真偽値の結果はイテレータの合成を容易にするために存在する
    • yield を適切にハンドリングしていれば、イテレータを受け取ってイテレータを返すような実装が簡単にできる
    • 途中で呼び出し側のイテレータの処理が中断されていても、他のイテレータも同様に処理が中断してくれる
  4. ループ本体内のすべての制御フロー文 (break、continue、defer、goto、return) の挙動は変化しない

だいぶ理解が深まった気がしてきたので、関連するパッケージについても軽く触れたいと思います！

iter パッケージについても見てみる

イテレータに関するパッケージとして iter パッケージが導入されました[1]。ソースコード自体は https://go.dev/src/iter/iter.go にあります。早速定義してある型や関数を見てみましょう！

型としては以下の2つが定義されているようです。

type Seq[V any] func(yield func(V) bool)

type Seq2[K, V any] func(yield func(K, V) bool)

3つのイテレータの形式のうち、func(yield func() bool) 以外の形式については、それぞれ Seq[V any] と Seq2[K, V any] として定義されるようです。

また、以下の2つの関数が実装されました。

func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func())

func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func())

早速、型と関数について深ぼってみます！

型を使ってイテレータの実装をしてみる

Seq[V any] を使ってみる！

Seq[V any] は func(yield func(V) bool) なので、先ほど登場したコードとほぼ同じですね。この型を使って、整数のスライスから奇数を取り除いてfor loopで処理するようなイテレータを実装してみます。

愚直に実装すれば以下のような実装になるかと思います。

func main() {
	for v := range FilterOutOdd {
		fmt.Println(v)
	}
}

func FilterOutOdd(yield func(int) bool) {
	for _, num := range []int{3, 1, 45, 91, 5, 2, 46, 9, 32, 534, 4, 10, 1} {
		if num%2 == 0 {
			if !yield(num) {
				return
			}
		}
	}
}

しかし、[]int{ ... } は FilterOutOdd に引数として渡したいケースの方が便利そうなので、以下の2点を修正します

  • []int{ ... } を FilterOutOdd に引数として渡す
  • FilterOutOdd の返り値を func(yield func(v int) bool) (つまり、 Seq[int] にする)

こんな感じでしょうか

func main() {
	nums := []int{3, 1, 45, 91, 5, 2, 46, 9, 32, 534, 4, 10, 1}
	for v := range FilterOutOdd(nums) {
		fmt.Println(v)
	}
}

func FilterOutOdd(nums []int) iter.Seq[int] {
	return func(yield func(int) bool) {
		for _, num := range nums {
			if num%2 == 0 {
				if !yield(num) {
					return
				}
			}
		}
	}
}

動かしてみてちゃんと意図した動作になっていることも確かめられました🎉
Seq[V any] を返り値として使うケースは実装イメージがわきやすいのではないでしょうか。

❯ go run .
2
46
32
534
4
10

Seq[K, V any] を使ってみる！

このケースも Seq[V any] とほぼ同じような実装ができそうです。
先ほどの FilterOutOdd の返り値に「4で割り切れるかどうか」というフラグを追加してみましょう。

  • FilterOutOdd の返り値を Seq2[int, bool] にする
  • yield に渡す引数を増やす

という2点を対応すれば良さそうですね！

実際のコードは以下のようになるでしょうか。

func main() {
	nums := []int{3, 1, 45, 91, 5, 2, 46, 9, 32, 534, 4, 10, 1}
	for v, canBeDividedBy4 := range FilterOutOdd(nums) {
		fmt.Printf("%dは4で割り切れるか？: %t\n", v, canBeDividedBy4)
	}
}

func FilterOutOdd(nums []int) iter.Seq2[int, bool] {
	return func(yield func(int, bool) bool) {
		for _, num := range nums {
			if num%2 == 0 {
				if !yield(num, num%4 == 0) {
					return
				}
			}
		}
	}
}

実行してみると、期待通りに動作していそうでした🎉
mapのkeyとvalueをイテレータで処理したい場合などは Seq[K, V any] を返り値とする関数を実装してあげれば良さそうですね。

❯ go run .
2は4で割り切れるか？: false
46は4で割り切れるか？: false
32は4で割り切れるか？: true
534は4で割り切れるか？: false
4は4で割り切れるか？: true
10は4で割り切れるか？: false

関数を使った実装を考えてみる！

続いて Pull 関数と Pull2 関数を見てみましょう。
関数のシグネチャだけをみても、どういう動作をするかイメージがつかみにくいので、コメントを読んでみます。

// Pull converts the “push-style” iterator sequence seq
// into a “pull-style” iterator accessed by the two functions
// next and stop.
//
// Next returns the next value in the sequence
// and a boolean indicating whether the value is valid.
// When the sequence is over, next returns the zero V and false.
// It is valid to call next after reaching the end of the sequence
// or after calling stop. These calls will continue
// to return the zero V and false.
//
// Stop ends the iteration. It must be called when the caller is
// no longer interested in next values and next has not yet
// signaled that the sequence is over (with a false boolean return).
// It is valid to call stop multiple times and when next has
// already returned false.
//
// It is an error to call next or stop from multiple goroutines
// simultaneously.

(意訳)

Pullは、「プッシュ型」反復子シーケンスseqを、2つの関数nextとstopによってアクセスされる「プル型」反復子に変換する。

nextはシーケンスの次の値と、その値が有効かどうかを示すブール値を返す。シーケンスが終了すると、nextはゼロのVとfalseを返す。シーケンスの終端に達した後やstopを呼び出した後にnextを呼び出すことは有効である。これらの呼び出しは、ゼロVとfalseを返し続ける。

stopは反復を終了する。stopは、呼び出し元がnextの値にもう興味がなく、nextがまだ（falseのブール値を返して）シーケンスの終了を通知していないときに呼び出されなければならない。stopを複数回呼び出しても、nextがすでにfalseを返していても有効である。

複数のゴルーチンから同時にnextやstopを呼び出すのはエラーです。

ざっくりまとめてみると Pull 関数は以下の役割や挙動をするようです。

  • Push型のイテレータをPull型に変換する
  • next はイテレータの次の値と、有効な値があるかどうかの真偽値を返す
  • stop はイテレータを終了させる
  • stop を呼び出した後、next からはゼロ値とfalseが常に返る

考え方としては、デザインパターンの「イテレータ」と似ていそうですね。Pull 関数のシグネチャは func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) なので、Goでは Seq[V any] の形式のイテレータをPush型と呼ぶようです。

先ほどの FilterOutOdd 関数を Pull 関数を用いて書き直してみます！

func main() {
	nums := []int{3, 1, 45, 91, 5, 2, 46, 9, 32, 534, 4, 10, 1}

	// イテレータを停止させる必要はないので、第2引数のstopは無視する
	next, _ := iter.Pull(FilterOutOdd(nums))
	for {
		v, ok := next()
		if !ok {
			break
		}
		fmt.Println(v)
	}
}

func FilterOutOdd(nums []int) iter.Seq[int] {
	return func(yield func(int) bool) {
		for _, num := range nums {
			if num%2 == 0 {
				if !yield(num) {
					return
				}
			}
		}
	}
}

for { ... } の無限ループ内で next 関数から値を受け取り、出力するようにしています。値がなくなったら ok がfalseを返すので、その場合に無限ループを抜け出すような実装になっています。

実際に動かしてみるとちゃんと期待通りの挙動になっていそうです🎉

❯ go run .
2
46
32
534
4
10

stop も使ってみましょう。

stop の後に next を呼び出してもゼロ値とfalseを返すとあるので、それを確かめてみます。v が50より大きければ、break を呼んで無限ループを抜けて、stop を呼び出した後、再び無限ループでイテレーションを続けてみます。

 	next, stop := iter.Pull(FilterOutOdd(nums))
 	for {
 		v, ok := next()
 		if !ok {
 			break
 		}
 		fmt.Println(v)
+		if v > 50 {
+			break
+		}
+	}
+
+	stop()
+
+	for {
+		v, ok := next()
+		if !ok {
+			break
+		}
+		fmt.Println(v)
	}

実行してみます。結果を見る限り、ドキュメントのきさstop の後に next を呼んでも値が取り出せていないので、ドキュメント通りの実装になっていそうです！

❯ go run .
2
46
32
534

stop の後に next を呼んでも値が取り出せていないので、ドキュメント通りの実装になっていそうです！

最後に、break で無限ループを抜けた後に、再度 next を呼び出してみます。

コードは以下のような感じで実装してみます。

+ 	next, _ := iter.Pull(FilterOutOdd(nums))
- 	next, stop := iter.Pull(FilterOutOdd(nums))
 	for {
 		v, ok := next()
 		if !ok {
 			break
 		}
 		fmt.Println(v)
		if v > 50 {
			break
		}
	}

+	fmt.Println("なんらかの処理")
-	stop()

	for {
		v, ok := next()
		if !ok {
			break
		}
		fmt.Println(v)
	}

実行してみると、以下のような結果になりました。

❯ go run .
2
46
32
534
なんらかの処理
4
10

Pull 関数を使えば、イテレータでの反復を呼び出し側で停止/再開できるようです。

個人的にはあまり使い所が思いつかないですが、Seq[V any] のようなPush型のイテレータだと反復を途中で止めて再開するということは基本的にできないと思われるので、そういったケースのみに有用な関数なのかなと思います。

イテレータを使ったGoのコードを読んでみる

ここまでイテレータの使い方と iter パッケージまで見てきました。イテレータについてかなり理解が深まったのではないでしょうか？

最後にイテレータを使ったGoの実装された関数について、詳細を見ていきたいと思います。

slices パッケージの Chunk 関数

実装コードは以下のようになっています。

func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
	if n < 1 {
		panic("cannot be less than 1")
	}

	return func(yield func(Slice) bool) {
		for i := 0; i < len(s); i += n {
			// Clamp the last chunk to the slice bound as necessary.
			end := min(n, len(s[i:]))

			// Set the capacity of each chunk so that appending to a chunk does
			// not modify the original slice.
			if !yield(s[i : i+end : i+end]) {
				return
			}
		}
	}
}

end はイテレータとして返すスライスの値をどこからどこまで取るかを示しているようですね。返り値は iter.Seq[Slice] なので、スライスが渡ってきそうです。

こんな感じでコードが書けます。

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	for chunk := range Chunk(nums, 3) {
		fmt.Println(chunk)
	}
}

実行してみます

❯ go run .
[1 2 3]
[4 5 6]
[7 8 9]
[10]

確かに長さ3のスライスがイテレータから渡ってきていそうですね！👏
過去に同じような実装を自前でやったことがある筆者にとっては、この上なく便利な関数になりそうです！！

slices パッケージには、そのほかにも Backward 関数など便利に使えそうな関数が実装されたので、興味のある人はソースコードを見てみると良いと思います！

maps パッケージの Values 関数

実装コードは以下のようになっています。

func Values[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[V] {
	return func(yield func(V) bool) {
		for _, v := range m {
			if !yield(v) {
				return
			}
		}
	}
}

実装はかなりシンプルですね！
ここまで呼んでくれた方ならその実装を読み解くのも容易なはずです！

こんな感じでコードが書けそうです。

func main() {
	maps := map[int]string{1: "one", 2: "two", 3: "three"}
	for v := range Values(maps) {
		fmt.Println(v)
	}
}

func Values[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[V] {
	return func(yield func(V) bool) {
		for _, v := range m {
			if !yield(v) {
				return
			}
		}
	}
}

実行してみます

❯ go run .
one
two
three

確かにvalueだけが取り出せていますね

x/exp/xiter パッケージの Limit 関数

先ほどの2つの例ではイテレータを使ってsliceやmapを扱いやすくするという関数でしたが、イテレータ自体に対する関数も x/exp/xiter パッケージで提供予定です[2]。

Limit 関数は以下のようなコードになっています。

func Limit[V any](seq Seq[V], n int) Seq[V] {
	return func(yield func(V) bool) {
		if n <= 0 {
			return
		}
		for v := range seq {
			if !yield(v) {
				return
			}
			if n--; n <= 0 {
				break
			}
		}
	}
}

引数がイテレータになっていて、n で指定された数だけそのイテレータを反復するという実装になっていそうですね。先ほどの FilterOutOdd 関数と組み合わせて、3つだけ偶数を取り出すという操作をしてみます。

コードはこんな感じになるでしょうか。

func main() {
	nums := []int{3, 1, 45, 91, 5, 2, 46, 9, 32, 534, 4, 10, 1}
	for n := range Limit(FilterOutOdd(nums), 3) {
		fmt.Println(n)
	}
}

func FilterOutOdd(nums []int) iter.Seq[int] {
	return func(yield func(int) bool) {
		for _, num := range nums {
			if num%2 == 0 {
				if !yield(num) {
					return
				}
			}
		}
	}
}

func Limit[V any](seq iter.Seq[V], n int) iter.Seq[V] {
	return func(yield func(V) bool) {
		if n <= 0 {
			return
		}
		for v := range seq {
			if !yield(v) {
				return
			}
			if n--; n <= 0 {
				break
			}
		}
	}
}

実行してみます

❯ go run .
2
46
32

確かに3つだけ偶数が出力されました🎉

まとめ

さて、みなさんいかがだったでしょうか？
イテレータに関する理解は深まったのではないかと思います！

個人的には yield からの返り値を適切に扱わないとランタイム時にパニックが起こる可能性があるというのが、実装上で気をつけなければならない点の1つかなと思いました。

とはいえ、イテレータがリリースされたことによって、Goの書き方の幅が広がったり、自前で実装する工数が減ることを考えると嬉しいことがたくさんありそうです！

みなさん良いイテレータライフを！！！ﾉｼ

脚注

  1. https://github.com/golang/go/issues/61897 ↩︎

  2. https://github.com/golang/go/issues/61898 ↩︎

Discussion