Zenn
🫰

Go の iter パッケージを使ってみよう

2024/08/22に公開
2件
Go
tech

はじめに

Go 1.23 で iter パッケージが導入されました。この iter は抽象化されたイテレータを示す仕組みと実装です。未だどの様に活用して良いか分からない方もいると思いますので、使い方を簡単に解説しようと思います。

概念

iter パッケージは、現状は for-range でのみ利用可能です。スコープにコンテキストを持ったロジカルな列挙可能オブジェクトと、それを別のスコープにて for-range でイテレートする際に便利です。
これまでであれば、こういった実装は goroutine と channel を使いスコープを分割させる事で実装してきました。

package main

func iter1[T any](a []T) func() (T, bool) {
	ch := make(chan T)
	go func() {
		defer close(ch)
		for _, v := range a {
			ch <- v
		}
	}()
	return func() (T, bool) {
		v, ok := <-ch
		return v, ok
	}
}

func main() {
	vv := iter1([]int{1, 2, 3})
	for {
		v, ok := vv()
		if !ok {
			break
		}
		println(v)
	}
}

しかしこのソースコードには問題があります。main 関数内で for を途中 break したい場合にも gorutine が残ってしまいます。これをちゃんと解決するには割と複雑なコードを書かないといけませんでした。

package main

import (
	"context"
)

func iter1[T any](ctx context.Context, a []T) func() (T, bool) {
	ch := make(chan T)
	go func() {
		defer close(ch)
		for _, v := range a {
			select {
			case ch <- v:
			case <-ctx.Done():
			}
		}
	}()
	return func() (T, bool) {
		v, ok := <-ch
		return v, ok
	}
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	vv := iter1(ctx, []int{1, 2, 3})
	for {
		v, ok := vv()
		if !ok {
			break
		}
		println(v)
		if v == 2 {
			cancel()
			break
		}
	}
}

またイテレートしたいという要件の為だけに goroutine や channel を用意するのは若干ですが無駄遣いです。

iter を使う

では iter を使うとどうなるのかを解説します。iter パッケージには2つの型が提供されています。

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

iter.Seq と iter.Seq2 の違いは for の使い方の違いです。slice を for で扱う場合の形式が iter.Seq

a := []int{1, 2, 3}
for v := range a {
    // using v
}

map を for で扱う場合の形式が iter.Seq2 と考えて下さい。

m := map[string]int{}
for k, v := range m {
    // using k or v
}

例えば iter.Seq の内容を全て画面に表示するのであれば以下の様に実装できます。

func PrintAll[V any](seq iter.Seq[V]) {
	for v := range seq {
		fmt.Println(v)
	}
}

こちらはこれまでの for-range に iter.Seq を渡しているだけなので、これまでの実装と何も変わりません。

iter を実装する

ちょっと難しいのは、iter の実装側です。例えばファイルを行で読み取り for-range に渡す例を示します。

package main

import (
	"bufio"
	"io"
	"iter"
	"log"
	"os"
)

func lines(r io.Reader) iter.Seq[string] {
	scanner := bufio.NewScanner(r)
	return func(yield func(string) bool) {
		for scanner.Scan() {
			if !yield(scanner.Text()) {
				break
			}
		}
	}
}

まず lines は行をイテレートさせるので、string の型を持った iter.Seq でなければなりません。関数のシグネチャは以下の様になります。

func lines(r io.Reader) iter.Seq[string]

使い方は以下の様になります。

func main() {
    f, err := os.Open("input.txt")
    if err != nil {
        log.Fatal(err)
    }
    defer f.Close()

    for line := range lines(f) {
        // using line
    }
}

引数は io.Reader を取るようにしましたが、ファイル名を取るでも構わないと思います。ただし関数 lines はファイルを開く際のエラーを返せる仕組みが無くなるため意図的に panic を起こすか以下の様な関数シグネチャを取らざるを得なくなります。

func lines(fname string) (iter.Seq[string], error) {
    /// 略
}

func main() {
    it, err := lines("input.txt")
    if err != nil {
        log.Fatal(err)
    }
    for line := range it {
        // using line
    }
}

結局のところ呼び出し側でエラーをハンドルする事になります。では lines の続きを見てみましょう。

func lines(r io.Reader) iter.Seq[string] {
	scanner := bufio.NewScanner(r)
	return func(yield func(string) bool) { // ※
		for scanner.Scan() {
			if !yield(scanner.Text()) {
				break
			}
		}
	}
}

func(yield func(string) bool) の部分を見て頂くと分かりますが、iter.Seq (iter.Seq2) の本体は関数です(※を参照)。この関数を1回呼び出す事と for-range のループを1回まわす事は等価になります。関数の引数に取る yield は、for の中身を実行する為のコルーチンです。引数には iter.Seq でイテレートされる値の型を持ち、戻り値の bool は false の場合に for 文が途中 break された事を示します。

よって yield 関数には、イテレートしたい型を渡し、その戻り値が false だった場合には即時に関数を抜けなければなりません。ここだけ覚えておけば問題ありません。

iter.Seq2 も考え方は同じす。例えば学校のクラスの生徒に対するテストの得点をロジカルにイテレートしたい場合、おおよそ以下の様な実装になると思います。

type ClassRoom struct {
    ...
}

func (cr *ClassRoom) Stores() iter.Seq2[string, int] {
    ...
}

func main() {
    cr, _ := loadClassRoom("3A")

    for name, score := range cr.Scores() {
        fmt.Printf("name=%v, score=%v\n", name, score)
    }
}

ここまでの例では単純なスライスを列挙していますが、実際には終了しないイテレータを作る事でエンドレスなループを実装する事もできます。

標準パッケージの iter

標準パッケージの幾らかでは既に iter が導入されています。代表的な物を示します。

slices

スライスを扱う際に便利な関数が追加されています。

func All[Slice ~[]E, E any](s Slice) iter.Seq2[int, E]
func Backward[Slice ~[]E, E any](s Slice) iter.Seq2[int, E]
func Values[Slice ~[]E, E any](s Slice) iter.Seq[E]
func AppendSeq[Slice ~[]E, E any](s Slice, seq iter.Seq[E]) Slice
func Collect[E any](seq iter.Seq[E]) []E
func Sorted[E cmp.Ordered](seq iter.Seq[E]) []E
func SortedFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E
func SortedStableFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E
func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice]

maps

マップを扱う際に便利な関数が追加されています。

func All[Map ~map[K]V, K comparable, V any](m Map) iter.Seq2[K, V]
func Keys[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[K]
func Values[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[V]
func Insert[Map ~map[K]V, K comparable, V any](m Map, seq iter.Seq2[K, V])
func Collect[K comparable, V any](seq iter.Seq2[K, V]) map[K]V

どちらのパッケージにも Collect が用意されている事から想像できる様に、iter によって得られる効果は、スコープの分離だけではなくメモリの節約もあるのです。

iter の効果

長い文字列を空白文字で分割し、それをイテレートする例を考えてみましょう。入力の文字列がどれだけ巨大かは実行してみないとわかりません。ですので strings.Split を使うとメモリが枯渇してしまうかもしれません。

package iter_test

import (
	"iter"
	"strings"
	"testing"
	"unicode/utf8"
)

const text = "Lorem ipsum dolor sit amet consectetur adipiscing elit. Vivamus cursus lorem eget vulputate vehicula. Vestibulum vitae tincidunt turpis egestas consectetur ligula. Nam non sapien lobortis viverra velit in pellentesque eros. Nam viverra purus eu iaculis vehicula. Nam vel ante urna. Suspendisse a vehicula libero a dapibus sem. In hac habitasse platea dictumst. Aliquam vulputate sagittis congue. Nam venenatis mauris in turpis volutpat sed dignissim libero consequat. Orci varius natoque penatibus et magnis dis parturient montes nascetur ridiculus mus. Integer eget consequat purus. Aliquam gravida ex vitae semper pretium magna sapien cursus lorem quis lacinia tellus metus eu ante. Interdum et malesuada fames ac ante ipsum primis in faucibus."

func BenchmarkWithoutIter(b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, s := range strings.Split(text, " ") {
			_ = s
		}
	}
}

func explodeSeq(s string) iter.Seq[string] {
	return func(yield func(string) bool) {
		for len(s) > 0 {
			_, size := utf8.DecodeRuneInString(s)
			if !yield(s[:size]) {
				return
			}
			s = s[size:]
		}
	}
}

func splitSeq(s, sep string, sepSave int) iter.Seq[string] {
	if len(sep) == 0 {
		return explodeSeq(s)
	}
	return func(yield func(string) bool) {
		for {
			i := strings.Index(s, sep)
			if i < 0 {
				break
			}
			frag := s[:i+sepSave]
			if !yield(frag) {
				return
			}
			s = s[i+len(sep):]
		}
		yield(s)
	}
}

func SplitSeq(s, sep string) iter.Seq[string] {
	return splitSeq(s, sep, 0)
}

func BenchmarkWithIter(b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for s := range SplitSeq(text, " ") {
			_ = s
		}
	}
}

strings.Split と iter.Seq を使ったベンチマークを用意しました。

goos: linux
goarch: amd64
pkg: iterbench
cpu: Intel(R) Core(TM) i5-8250U CPU @ 1.60GHz
BenchmarkWithoutIter-8   	 190472	     5547 ns/op	   1792 B/op	      1 allocs/op
BenchmarkWithIter-8      	 492627	     2129 ns/op	     24 B/op	      2 allocs/op
PASS
ok  	iterbench	2.743s

1オペレーションあたりでは iter の方がメモリ確保が1回多いですが、実行処理時間やメモリ使用量では大きな違いが生じています。

可能であれば iter パッケージを導入して頑丈でかつメンテしやすい実装を目指して下さい。

おわりに

Go 1.23 に導入された iter パッケージについて解説しました。関数シグネチャが若干複雑なのでとっつきにくいかもしれませんが、これらはほぼ定形です。1度実装してしまえばイテレートする型を変えて参考にするだけなので意外と障壁は高くありません。ぜひ試してみると良いと思います。

以下のリポジトリに for-range-experiment と iter を使った例をいくつか用意しています。参考にして頂ければと思います。

https://github.com/mattn/go-for-range-experiment-example

Discussion