迭代器

寒江蓑笠翁大约 30 分钟

迭代器

在Go中,用于迭代特定数据结构的关键字为for range,之前的章节中已经介绍过它的一些应用,它仅能作用于语言内置的几个数据结构

  • 数组
  • 切片
  • 字符串
  • map
  • chan
  • 整型值

这样的话使用起来非常的不灵活,没有拓展性,对于自定义类型几乎不支持,不过好在go1.23版本更新以后,for range关键字支持了range over func,这样一来自定义迭代器也就成为了可能。

认识

下面通过一个例子来初步认识迭代器,不知道各位还是否记得在函数小节中讲解的闭包求解斐波那契数列的例子,它的实现代码如下

func Fibonacci(n int) func() (int, bool) {
	a, b, c := 1, 1, 2
	i := 0
	return func() (int, bool) {
		if i >= n {
			return 0, false
		} else if i < 2 {
			f := i
			i++
			return f, true
		}

		a, b = b, c
		c = a + b
		i++

		return a, true
	}
}

我们可以把它改造成迭代器,如下所示,可以看到代码量要减少了一些

func Fibonacci(n int) func(yield func(int) bool) {
	a, b, c := 0, 1, 1
	return func(yield func(int) bool) {
		for range n {
			if !yield(a) {
				return
			}
			a, b = b, c
			c = a + b
		}
	}
}

Go的迭代器是range over func风格,我们可以直接用for range关键字来进行使用,使用起来也要比原来更方便

func main() {
    n := 8
	for f := range Fibonacci(n) {
		fmt.Println(f)
	}
}

输出如下

0
1
1
2
3
5
8
13

如上所示,迭代器就是一个闭包函数,它接受一个回调函数作为参数,你甚至可以在里面看到yield这种字眼,写过python的人应该都很熟悉,它与python中的生成器很类似。Go的迭代器并没有新增任何关键字,语法特性,在上述示例中yield也只是一个回调函数,它并非关键字,官方取这个名字是为了方便理解。

推送式迭代器

关于迭代器的定义,我们可以在iter库中找到如下解释

An iterator is a function that passes successive elements of a sequence to a callback function, conventionally named yield.

迭代器是一个函数,它将序列中的元素逐个传递给回调函数,通常称为 yield。

我们从中可以明确的一点,迭代器就是一个函数,它接受一个回调函数作为参数,在迭代过程中会将序列中的元素逐个传递给回调函数yield。在之前示例中我们是按照下面的方式使用迭代器的

for f := range Fibonacci(n) {
    fmt.Println(f)
}

根据官方定义,上面迭代器Backward的例子使用就等同于下面这段代码

Fibonacci(n)(func(f int) bool {
    fmt.Println(f)
    return true
})

循环体的body就是迭代器的回调函数yiled,当函数返回true迭代器会继续迭代,否则就会停止。

此外,iter标准库中也定义了迭代器的类型iter.Seq,它的类型就是函数。

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

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

iter.Seq的回调函数只接受一个参数,那么在迭代时for range仅有一个返回值,如下

for v := range iter {
	// body
}

iter.Seq2的回调函数接受两个参数,那么在迭代时for range就有两个返回值,如下

for k, v := range iter {
	// body
}

虽然标准库中没有定义0个参数的Seq,但这也是完全允许的,它相当于

func(yield func() bool)

使用起来如下所示

for range iter {
	// body
}

回调函数的参数数量只能是0至2个,多了会无法通过编译。

简而言之,for range中的循环体就是迭代器中的yield回调函数,for range返回几个值,相应的yeild函数就有几个入参,每一轮迭代时,迭代器都会调用yield函数,也就是执行循环体中的代码,主动将序列中的元素传递给yield函数,这种主动传递元素的迭代器我们一般称之为推送式迭代器(pushing iterator),比较典型的例子就是其他语言中的foreach,比如js

let arr = [1, 2, 3, 4, 5];
arr.filter(e => e % 2 === 0).forEach(e => {
        console.log(e)
    });

在Go中的表现形式就是由range返回被迭代的元素。

for index, value := range iterator() {
	fmt.Println(index, value)
}

在某些语言(比如Java)中它还有另一个叫法:数据流处理。

既然循环体中的代码是作为回调函数传入迭代器的,而且它很可能是一个闭包函数,Go就需要让一个闭包函数在执行deferreturnbreakgoto等关键字时表现的像一个普通循环体代码段一样,思考下面几种情况。

比如说在迭代器循环中返回,那么在yield回调函数中要怎么去处理这个return呢?

for index, value := range iterator() {
    if value > 10 {
        return
	}
	fmt.Println(index, value)
}

不可能直接在回调函数中return,这么做只会让迭代停止而已,达不到返回的效果

iterator()(func(index int, value int) bool {
	if value > 10 {
		return false
	}
	fmt.Println(index, value)
})

再比如说在迭代器循环中使用defer

for index, value := range iterator() {
    defer fmt.Println(index, value) 
}

也不能直接在回调函数中使用defer,因为这么做的话在回调函数结束时就会直接延迟调用了

iterator()(func(index int, value int) bool {
	defer fmt.Println(index, value) 
})

像其他的几个关键字breakcontinuegoto也是类似的,好在这些情况Go已经帮我们处理好了,我们只需使用即可,可以暂时不需要关心这些,如果感兴趣可以自行浏览rangefunc/rewrite.goopen in new window中的源代码。


拉取式迭代器

推送式迭代器(pushing iterator)是由迭代器来控制迭代的逻辑,用户被动获取元素,相反的拉取式迭代器(pulling iterator)就是由用户来控制迭代逻辑,主动的去获取序列元素。一般而言,拉取式迭代器都会有特定的函数如next()stop()来控制迭代的开始或结束,它可以是一个闭包或者结构体。

scanner := bufio.NewScanner(file)
for scanner.Scan() {
    line, err := scanner.Text(), scanner.Err()
    if err != nil {
        fmt.Println(err)
        return
    }
    fmt.Println(line)
}

如上所示,Scanner通过方法Text()来获取文件中的下一行文本,通过方法Scan()来表示迭代是否结束,这也是拉取式迭代器的一种模式。Scanner采用结构体来记录状态,而在iter库定义的拉取式迭代器采用闭包来记录状态,我们通过iter.Pulliter.Pull2函数就可以将一个标准的推送式迭代器转换为拉取式迭代器,iter.Pulliter.Pull2的区别就是后者的返回值有两个,签名如下

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())

它们都接受一个迭代器作为参数,然后会返回两个函数next()stop(),用于控制迭代的继续和停止。

func next() (V, bool)

func stop()

next会返回被迭代的元素,和一个表示当前值是否有效的布尔值,当迭代结束时next函数会返回元素的零值和falsestop函数会结束迭代过程,当调用者不再使用迭代器后,就必须使用stop函数来结束迭代。顺带一提,在多个协程调用同一个迭代器的next函数是错误的做法,因为它并非并发安全。


下面通过一个例子来演示,它的功能就是把之前的斐波那契迭代器改造成拉取式迭代器,如下

func main() {
	n := 10
	next, stop := iter.Pull(Fibonacci(n))
	defer stop()
	for {
		fibn, ok := next()
		if !ok {
			break
		}
		fmt.Println(fibn)
	}
}

输出

0
1
1
2
3
5
8
13
21
34

这样一来我们就可以通过nextstop函数来手动控制迭代的逻辑了。你或许可能会觉得这样做多此一举,如果要这样做的话那为什么不直接用最初的闭包版本就好了,一样可以自己控制迭代,闭包的用法是这样的

func main() {
	fib := Fibonacci(10)
    for {
        n, ok := fib()
        if !ok {
            break
        }
        fmt.Prinlnt(n)
    }
}

转换过程:闭包 → 迭代器 → 拉取式迭代器,闭包与拉取式迭代器的用法都大差不差,它们的思想都是一样的,后者还会因为各种各样的处理导致性能上的拖累。老实说这么做确实多此一举,它的应用场景确实不是很多,不过iter.pull是为了iter.Seq而存在的,也就是为了将推送式迭代器转换成拉取式迭代器的而存在的,如果你仅仅只是想要一个拉取式迭代器,还专门为此去实现一个推送式迭代器来进行转换,要这样做的话不妨考虑下自己实现的复杂度和性能,就像这个斐波那契数列的例子一样,绕了一圈又回到原点,唯一的好处可能就是符合官方的迭代器规范。

错误处理

在迭代时发生了错误怎么办?我们可以将其传递给yield函数让for range返回,让调用者来进行处理,就像下面这个行迭代器的例子一样

func ScanLines(reader io.Reader) iter.Seq2[string, error] {
	scanner := bufio.NewScanner(reader)
	return func(yield func(string, error) bool) {
		for scanner.Scan() {
			if !yield(scanner.Text(), scanner.Err()) {
				return
			}
		}
	}
}

提示

值得注意的是,ScanLines迭代器是一次性使用的,文件关闭以后就不能再次使用了。

可以看到它的第二个返回值是error类型,使用起来如下

for line, err := range ScanLines(file) {
    if err != nil {
        fmt.Println(err)
        break
    }
    fmt.Println(line)
}

这样处理起来就跟普通的错误处理没什么区别,拉取式迭代器也是同理

next, stop := iter.Pull2(ScanLines(file))
defer stop()
for {
    line, err, ok := next()
    if err != nil {
        fmt.Println(err)
        break
    } else if !ok {
        break
    }
    fmt.Println(line)
}

如果发生了panic,就像平常一样使用recovery即可。

defer func() {
    if err := recover(); err != nil {
        fmt.Println("panic:", err)
        os.Exit(1)
    }
}()

for line, err := range ScanLines(file) {
    if err != nil {
        fmt.Println(err)
        break
    }
    fmt.Println(line)
}

拉取式迭代器依然同理,这里就不演示了。

标准库

有很多标准库也支持了迭代器,最常用的就是slicesmaps标准库,下面介绍几个比较实用的功能。

slices.All

func All[Slice ~[]E, E any](s Slice) iter.Seq2[int, E] 

slices.All会将切片转换成一个切片迭代器

func main() {
	s := []int{1, 2, 3, 4, 5}
	for i, n := range slices.All(s) {
		fmt.Println(i, n)
	}
}

输出

0 1
1 2
2 3
3 4
4 5

slices.Values

func Values[Slice ~[]E, E any](s Slice) iter.Seq[E]

slices.Values会将切片转换成一个切片迭代器,但是不带索引

func main() {
	s := []int{1, 2, 3, 4, 5}
	for n := range slices.Values(s) {
		fmt.Println(n)
	}
}

输出

1
2
3
4
5

slices.Chunk

func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice]

slices.Chunk函数会返回一个迭代器,该迭代器会以n个元素为切片推送给调用者

func main() {
	s := []int{1, 2, 3, 4, 5}
	for chunk := range slices.Chunk(s, 2) {
		fmt.Println(chunk)
	}
}

输出

[1 2]
[3 4]
[5]

slices.Collect

func Collect[E any](seq iter.Seq[E]) []E

slices.Collect函数会将切片迭代器收集成一个切片

func main() {
	s := []int{1, 2, 3, 4, 5}
	s2 := slices.Collect(slices.Values(s))
	fmt.Println(s2)
}

输出

[1 2 3 4 5]

maps.Keys

func Keys[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[K]

maps.Keys会返回一个迭代map所有键的迭代器,配合slices.Collect可以直接收集成一个切片。

func main() {
	m := map[string]int{"one": 1, "two": 2, "three": 3}
	keys := slices.Collect(maps.Keys(m))
	fmt.Println(keys)
}

输出

[three one two]

由于map是无序的,所以输出也不固定

maps.Values

func Values[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[V]

maps.Values会返回一个迭代map所有值的迭代器,配合slices.Collect可以直接收集成一个切片。

func main() {
	m := map[string]int{"one": 1, "two": 2, "three": 3}
	keys := slices.Collect(maps.Values(m))
	fmt.Println(keys)
}

输出

[3 1 2]

由于map是无序的,所以输出也不固定

maps.All

func All[Map ~map[K]V, K comparable, V any](m Map) iter.Seq2[K, V]

maps.All可以将一个map转换为成一个map迭代器

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

一般不会这么直接用,都是拿来配合其他数据流处理函数的。

maps.Collect

func Collect[K comparable, V any](seq iter.Seq2[K, V]) map[K]V

maps.Collect可以将一个map迭代器收集成一个map

func main() {
	m := map[string]int{"one": 1, "two": 2, "three": 3}
	m2 := maps.Collect(maps.All(m))
	fmt.Println(m2)
}

collect函数一般作为数据流处理的终结函数来使用。

链式调用

通过上面标准库提供的函数,我们可以将其组合来处理数据流,比如对数据流进行排序,如下

sortedSlices := slices.Sorted(slices.Values(s))

go的迭代器采用的是闭包,只能像这样嵌套函数调用,本身没法链式调用,调用链长了以后可读性会很差,但我们可以自己通过结构体来记录迭代器,就能够实现链式调用。

demo

一个简单的链式调用demo如下所示,它包含了FilterMapFindSome等常用的功能。

package iterx

import (
	"iter"
	"slices"
)

type SliceSeq[E any] struct {
	seq iter.Seq2[int, E]
}

func (s SliceSeq[E]) All() iter.Seq2[int, E] {
	return s.seq
}

func (s SliceSeq[E]) Filter(filter func(int, E) bool) SliceSeq[E] {
	return SliceSeq[E]{
		seq: func(yield func(int, E) bool) {
			// 重新组织索引
			i := 0
			for k, v := range s.seq {
				if filter(k, v) {
					if !yield(i, v) {
						return
					}
					i++
				}
			}
		},
	}
}

func (s SliceSeq[E]) Map(mapFn func(E) E) SliceSeq[E] {
	return SliceSeq[E]{
		seq: func(yield func(int, E) bool) {
			for k, v := range s.seq {
				if !yield(k, mapFn(v)) {
					return
				}
			}
		},
	}
}

func (s SliceSeq[E]) Fill(fill E) SliceSeq[E] {
	return SliceSeq[E]{
		seq: func(yield func(int, E) bool) {
			for i, _ := range s.seq {
				if !yield(i, fill) {
					return
				}
			}
		},
	}
}

func (s SliceSeq[E]) Find(equal func(int, E) bool) (_ E) {
	for i, v := range s.seq {
		if equal(i, v) {
			return v
		}
	}
	return
}

func (s SliceSeq[E]) Some(match func(int, E) bool) bool {
	for i, v := range s.seq {
		if match(i, v) {
			return true
		}
	}
	return false
}

func (s SliceSeq[E]) Every(match func(int, E) bool) bool {
	for i, v := range s.seq {
		if !match(i, v) {
			return false
		}
	}
	return true
}

func (s SliceSeq[E]) Collect() []E {
	var res []E
	for _, v := range s.seq {
		res = append(res, v)
	}
	return res
}

func (s SliceSeq[E]) Sort(cmp func(x, y E) int) []E {
	collect := s.Collect()
	slices.SortFunc(collect, cmp)
	return collect
}

func (s SliceSeq[E]) SortStable(cmp func(x, y E) int) []E {
	collect := s.Collect()
	slices.SortStableFunc(collect, cmp)
	return collect
}

func Slice[S ~[]E, E any](s S) SliceSeq[E] {
	return SliceSeq[E]{seq: slices.All(s)}
}

然后我们就可以通过链式调用来处理了,看几个使用案例。

处理元素值

func main() {
	s := []string{"apple", "banana", "cherry"}
	all := iterx.Slice(s).Map(strings.ToUpper).All()
	for i, v := range all {
		fmt.Printf("index: %d, value: %s\n", i, v)
	}
}

输出

index: 0, value: APPLE
index: 1, value: BANANA
index: 2, value: CHERRY

寻找某一个指定值

func main() {
	s := []int{1, 2, 3, 4, 5}
	result := iterx.Slice(s).Find(func(i int, e int) bool {
		return e == 3
	})
	fmt.Println(result)
}

输出

3

填充切片

func main() {
	s := []int{1, 2, 3, 4, 5}
	result := iterx.Slice(s).Fill(6).Collect()
	fmt.Println(result)
}

输出

[6 6 6 6 6]

过滤元素

func main() {
	s := []int{1, 2, 3, 4, 5}
	filter := iterx.Slice(s).Filter(func(i int, e int) bool {
		return e%2 == 0
	}).All()
	for i, v := range filter {
		fmt.Printf("Index: %d, Value: %d\n", i, v)
	}
}

输出

Index: 0, Value: 2
Index: 1, Value: 4

比较可惜的是Go目前还不支持简写匿名函数,就像js,rust,java中的箭头函数一样,否则链式调用还可以更加简洁和优雅一些。

性能

因为Go对迭代器做了许多的处理,它的性能肯定是不如原生for range循环的,我们拿最简单的一个切片遍历来测试下它们的性能区别,分为下面几种

  • 原生for循环
  • 推送式迭代器
  • 拉取式迭代器

测试代码如下,测试切片长度为1000。

package main

import (
	"iter"
	"slices"
	"testing"
)

var s []int

const n = 10000

func init() {
	for i := range n {
		s = append(s, i)
	}
}

func testNaiveFor(s []int) {
	for i, n := range s {
		_ = i
		_ = n
	}
}

func testPushing(s []int) {
	for i, n := range slices.All(s) {
		_ = i
		_ = n
	}
}

func testPulling(s []int) {
	next, stop := iter.Pull2(slices.All(s))
	for {
		i, n, ok := next()
		if !ok {
			stop()
			return
		}
		_ = i
		_ = n
	}
}

func BenchmarkNaive_10000(b *testing.B) {
	for range b.N {
		testNaiveFor(s)
	}
}

func BenchmarkPushing_10000(b *testing.B) {
	for range b.N {
		testPushing(s)
	}
}

func BenchmarkPulling_10000(b *testing.B) {
	for range b.N {
		testPulling(s)
	}
}

测试结果如下

goos: windows
goarch: amd64
pkg: golearn
cpu: 11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz
BenchmarkNaive_10000
BenchmarkNaive_10000-16           492658              2398 ns/op               0 B/op          0 allocs/op
BenchmarkPushing_10000
BenchmarkPushing_10000-16         315889              3707 ns/op               0 B/op          0 allocs/op
BenchmarkPulling_10000
BenchmarkPulling_10000-16           2016            574509 ns/op             440 B/op         14 allocs/op
PASS
ok      golearn 4.029s

我们通过结果可以看到推送式迭代器与原生的for range循环相差不是特别大,但拉取式迭代器要比前面两个慢了几乎两个数量级,在使用的时候各位可以根据自己的实际情况来进行考虑。

小结

与泛型的情况相似,Go的迭代器同样饱受争议,部分人的观点是迭代器引入了过多的复杂度,违背了Go的简洁哲学,像这种迭代器的闭包代码多了以后,调试起来怕是都有点困难,阅读起来就就更加恼火了。

你可以在很多地方看到关于迭代器的激烈讨论

理性的看待Go迭代器,它确实使得编写代码更加方便,尤其是在处理切片类型的时候,但同时也会引入了些许复杂度,迭代器部分的代码可读性会降低,不过总的来说,我认为这确实是一个实用的特性。