SODA Engineering Blog
🔁

Go1.23のイテレータを動かして学ぶ

2024/06/23に公開

はじめに(おことわり)

Goの次期バージョンである1.23でイテレータが導入されるようなので触ってみます。
今回の調査で使うGoのバージョンは1.23rc1です。公式ドキュメントからの情報に基づいていますが、RC版ですので、正式リリース時には変更が入っている可能性がありますので、注意が必要です。

イテレータとは

イテレータは、連続したデータを1つずつ順番に処理していくための仕組みです。
イテレータをを使うと、スライスを使わずにシーケンシャルなデータを扱うことができます。

Goのイテレータは、コールバック関数を受け取る関数です。返り値はありません。コールバック関数を呼び出してデータを1つずつ渡していきます。
Goのイテレータ関数は次のようなシグネチャで定義されます。コールバック関数は慣習としてyieldと名付けられます。
コールバック関数は引数を1つ受け取るものと、2つ受け取るものがあります。KVは、マップのKeyとValueとなるような値をとります。

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

また、Goのイテレータには、PushスタイルとPullスタイルがあります。
Pushスタイルは1つずつデータを渡していきます。逆にPullスタイルは1つずつデータを受け取ります。

Pushスタイルのイテレータを触ってみる

シンプルな例

使い方は、イテレータ関数を定義してforrangeに渡すだけです。このrangeに関数が渡せるようになること自体も、今後導入が予定されている新機能です。range-over functionと呼ばれています。
今回は、intの値を受け取るコールバック関数func (yield func(int) bool)を使って説明します。使い方はfunc (yield func(K, V) bool)の場合でも同じです。

package main

import "fmt"

// イテレータ関数を定義する
func iterator(yield func(int) bool) {
	// コールバック関数に値をPushしていく
	yield(1)
	yield(2)
	yield(3)
}

func main() {
	// Pushされた値はループ変数で受け取ることができる
	for i := range iterator {
		fmt.Println(i)
	}
}

これを実行すると次の結果が出力されます。

1
2
3

イテレータ関数の中でコールバック関数を呼び出して、値を渡していきます。このとき、コールバック関数yieldを呼び出すたびにイテレータ関数の実行が一時停止し、制御がforループに移ります。そして、ループ処理を終えると、またイテレータ関数に制御が戻り、一時停止していたyieldの呼び出し位置から処理が再開されます。
yieldの呼び出しの間にログを仕込むと処理の流れがわかりやすくなります。

package main

import "fmt"

func iterator(yield func(int) bool) {
	yield(1)
	fmt.Println("after yield(1)")
	yield(2)
	fmt.Println("after yield(2)")
	yield(3)
	fmt.Println("after yield(3)")
}

func main() {
	for i := range iterator {
		fmt.Println(i)
	}
}
1
after yield(1)
2
after yield(2)
3
after yield(3)

イテレータ関数とループを行ったり来たりしていることがわかります。

イテレータを繋げてみる

イテレータを引数で受け取ってイテレータを返す関数を作ることで、イテレータの処理を繋げることができます。
ここでは、0から9までのシーケンシャルなデータの中から偶数の値だけを取得する処理と、それを2倍にする処理をイテレータを使って繋げてみます。
rangeintを渡してしますが、これはGo1.22から導入されたrange-over intという機能を使っています。rangeにintの値Nを渡すと、0からN-1の整数を連続で返してくれます。

package main

import (
	"fmt"
)

// 0から9の整数をPushするイテレータを返す関数
func numbers() func(func(int) bool) {
	return func(yield func(int) bool) {
		// range-over intで生成された整数をPushしている
		for i := range 10 {
			yield(i)
		}
	}
}

// イテレータを受け取って、偶数の場合だけPushするイテレータを返す関数
func even(seq func(func(int) bool)) func(func(int) bool) {
	// 偶数の場合だけPushするイテレータ
	return func(yield func(int) bool) {
		// 関数外から渡されたイテレータをrangeに渡して値をもらう
		for i := range seq {
			if i%2 == 0 {
				yield(i)
			}
		}
	}
}

// イテレータを受け取って、要素を2倍にしてPushするイテレータを返す関数
func double(seq func(func(int) bool)) func(func(int) bool) {
	// 各要素を2倍にしてPushするイテレータ
	return func(yield func(int) bool) {
		// 関数外から渡されたイテレータをrangeに渡して値をもらう
		for i := range seq {
			yield(i * 2)
		}
	}
}

func main() {
	// numbers内の偶数だけが2倍になって取得される
	for i := range double(even(numbers())) {
		fmt.Println(i)
	}
}

実行すると大元のイテレータが生成する値のうち、偶数だけが2倍になってループ変数に渡ってきていることがわかります。

0
4
8
12
16

便利なiter.Seqとiter.Seq2

標準ライブラリのiterに定義されているSeqSeq2を使うと、よりわかりやすくイテレータを書くことができます。
SeqSeq2はイテレータのシグネチャを持つ関数として定義されています。
Seqはsequenceの略です。

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

先の例をiter.Seqを用いて書き直すと次のようになります。
numbers関数の返り値、even関数とdouble関数の引数と返り値がiter.Seqに置き換わっています。実行結果は変わりません。

package main

import (
	"fmt"
	"iter"
)

// 0から9の整数をPushするイテレータを返す関数
func numbers() iter.Seq[int] {
	return func(yield func(int) bool) {
		// range-over intで生成された整数をPushしている
		for i := range 10 {
			yield(i)
		}
	}
}

// イテレータを受け取って、偶数の場合だけPushするイテレータを返す関数
func even(seq iter.Seq[int]) iter.Seq[int] {
	// 偶数の場合だけPushするイテレータ
	return func(yield func(int) bool) {
		// 関数外から渡されたイテレータをrangeに渡して値をもらう
		for i := range seq {
			if i%2 == 0 {
				yield(i)
			}
		}
	}
}

// イテレータを受け取って、要素を2倍にしてPushするイテレータを返す関数
func double(seq iter.Seq[int]) iter.Seq[int] {
	// 各要素を2倍にしてPushするイテレータ
	return func(yield func(int) bool) {
		// 関数外から渡されたイテレータをrangeに渡して値をもらう
		for i := range seq {
			yield(i * 2)
		}
	}
}

func main() {
	// numbers内の偶数だけが2倍になって取得される
	for i := range double(even(numbers())) {
		fmt.Println(i)
	}
}

Pullスタイルのイテレータを触ってみる

これまでみてきた通り、range-over functionを使えばイテレータからPushされた値を受け取ることができました。しかし、PushスタイルのイテレータをPullスタイルに変換することでも値を受けとることができます。

先の例で使ったnumbers関数が返すイテレータをPullスタイルに変換して使ってみます。
PushスタイルのイテレータをPullスタイルに変換するにはiter.Pull関数を使用します。
iter.Pull関数はnextstopの2つの関数を返します。
next関数はシーケンシャルなデータから値を1つずつ順番に取得します。また、次の要素が残っているかを表すフラグをboolで返します。残っている場合はtrue、そうではない場合はfalseが返ってきます。
stop関数はイテレーションを任意の場所で止めるための関数です。これ以上イテレーションが不要になった場合は、必ずstop関数を呼び出す必要があります。

package main

import (
	"fmt"
	"iter"
)

// 0から9の整数をPushするイテレータを返す関数
func numbers() iter.Seq[int] {
	return func(yield func(int) bool) {
		// range-over intで生成された整数をPushしている
		for i := range 10 {
			yield(i)
		}
	}
}

func main() {
	seq := numbers()
	next, stop := iter.Pull(seq)
	defer stop()
	for {
		i, ok := next()
		if !ok {
			break
		}
		fmt.Println(i)
	}
}

実行結果は次のとおりです。

0
1
2
3
4
5
6
7
8
9

イテレータを使うときの注意点

イテレーションの止め方(range-over function編)

Pushスタイルのイテレータをrange-over functionで使用した場合、イテレータの内部で呼び出すコールバック関数yieldの返り値をハンドリングする必要があります。
range-over functionを使ったループをbreakで抜ける場合、コールバック関数yieldの返り値はfalseになります。falseが返ってきた場合はこれ以上イテレーションをする必要がないので、それ以降のyield関数の呼び出しを止めます。
イテレーションを止める必要がある場合にyield関数を呼び出すと、実行時エラーになるので注意が必要です。

package main

import (
	"iter"
	"fmt"
)

// 0から9の整数をPushするイテレータを返す関数
func numbers() iter.Seq[int] {
	return func(yield func(int) bool) {
		for i := range 10 {
			// 返り値を適切にハンドリングしていないのでpanicになる可能性がある
			yield(i)
		}
	}
}

func main() {
	seq := numbers()
	for i := range seq {
		fmt.Println(i)
		// breakしてイテレーションから抜ける
		break
	}
}
panic: runtime error: range function continued iteration after function for loop body returned false

yield関数の返り値をハンドリングして、これ以上イテレーションが必要ない場合はイテレータ関数内でreturnしてイテレーションを止めます。これでエラーにならずイテレーションを止めることができます。

package main

import (
	"iter"
	"fmt"
)

// 0から9の整数をPushするイテレータを返す関数
func numbers() iter.Seq[int] {
	return func(yield func(int) bool) {
		for i := range 10 {
			// 返り値を適切にハンドリングしていないのでpanicになる可能性がある
			if !yield(i) {
				// returnしてイテレーションを止める
				return
			}
		}
	}
}

func main() {
	seq := numbers()
	for i := range seq {
		fmt.Println(i)
		// breakしてイテレーションから抜ける
		break
	}
}

イテレーションの止め方(Pullスタイルイテレータ編)

Pullスタイルのイテレータでは、next関数の返り値をハンドリングして、これ以上のイテレーションが必要かを判断します。
Pullスタイルのイテレータでは条件式なしのforを使うことが想定されますが、next関数の2番目の返り値をハンドリングしてbreakしないと無限ループになってしまいます。
ループ内でstop関数を呼び出すとイテレータはシーケンシャルなデータの値を返さなくなりますが、ループの実行自体は止まりません。next関数は1番目の返り値としてゼロ値を返し続けます。

package main

import (
	"fmt"
	"iter"
)

func number() iter.Seq[int] {
	return func(yield func(int) bool) {
		// 100を1度だけPush
		if !yield(100) {
			return
		}
	}
}

func main() {
	seq := number()
	next, stop := iter.Pull(seq)
	for {
		i, _ := next()
		fmt.Println(i)
		stop()
	}
}

正しくイテレーションを止めるには、next関数の2番目の返り値をハンドリングします。

package main

import (
	"fmt"
	"iter"
)

func number() iter.Seq[int] {
	return func(yield func(int) bool) {
		// 100を1度だけPush
		if !yield(100) {
			return
		}
	}
}

func main() {
	seq := number()
	next, stop := iter.Pull(seq)
	defer stop()
	for {
		i, ok := next()
		if !ok {
			// これ以上イテレーションの必要がないのでループを抜ける
			break
		}
		fmt.Println(i)
	}
}

一度停止したイテレータは再度実行できない

一度停止したイテレータは再度使用することはできません。シーケンスを最後まで読み取った場合やイテレーションの途中でstopしたイテレータを再度実行するとゼロ値が返ってきます。

package main

import (
	"fmt"
	"iter"
)

func numbers() iter.Seq[int] {
	return func(yield func(int) bool) {
		for i := range 10 {
			if !yield(i) {
				return
			}
		}
	}
}

func main() {
	seq := numbers()
	next, stop := iter.Pull(seq)

	i1, ok1 := next() // 有効
	i2, ok2 := next() // 有効
	stop()
	i3, ok3 := next() // 無効

	fmt.Printf("i1, ok1 = %d, %v\n", i1, ok1)
	fmt.Printf("i2, ok2 = %d, %v\n", i2, ok2)
	fmt.Printf("i3, ok3 = %d, %v\n", i3, ok3)
}

stop関数呼び出し以降のnext関数は無効な値を返していることがわかります。

i1, ok1 = 0, true
i2, ok2 = 1, true
i3, ok3 = 0, false

おわりに

冒頭でも述べましたが、今回の調査で使ったGoのバージョンは1.23rc1です。正式リリース時には変更が入っている可能性があるので、リリース後にプロダクションで使用する場合は公式ドキュメントを確認してください。
また、今回はイテレータのコールバック関数で引数を2つ受け取るfunc (yield func(K, V) bool)のイテレータのサンプルはありませんでした。しかし、使い方は引数1つのものと変わらないので、気になる方は参考資料にある公式ドキュメントをみて試してみてください。

参考資料

SODA Engineering Blog
SODA Engineering Blog

Discussion