跳转至

浅谈 Go 1.23 迭代器及 iter 包

很多流行的编程语言中都以某种方式提供迭代器，其中包括 C++、Java、Javascript、Python 和 Rust。Go 语言现在也加入了迭代器。iter 包是 Go 1.23 新增的标准库，提供了迭代器的基本定义和相关操作。

为什么需要迭代器

在 Go 1.18 引入泛型之后，便可以很方便的定义一些泛型容器类型来提升编码效率。我们可以定义很多不同的容器类型。而对于容器类型，通常都会需要遍历其中的所有元素。但是如何实现遍历功能，不同的容器类型、不同的人都会有不同的实现方案。

当我们拿到一个别人定义好的容器类型，我们必须先了解一下容器的内部细节，然后才能去尝试遍历它，但这样就无法编写一个适用于多种不同类型容器的迭代函数。因此， Go 生态系统需要有一个迭代容器内元素的官方标准，让 Gopher 们能够尽量统一起来。

从 Go 1.23 开始，for/range 语句支持单参数的函数。 并且单参数本身必须是一个接收 0 到 2 个参数并返回一个 bool 的函数；按照惯例，我们称之为 yield 函数。

Go
1
2
3
func(yield func() bool)
func(yield func(V) bool)
func(yield func(K, V) bool)

当我们在 Go 中谈论迭代器时，我们指的是具有这三种类型之一的函数。

Go 迭代器定义

Go 1.23 中增加的迭代器是一个函数类型，它把另一个函数（ yield 函数）作为参数，将容器中的连续元素传递给 yield 函数。

迭代器函数会在以下两种情况停止。

  • 在序列迭代结束后停止，表示此次迭代结束。
  • 在 yield 函数返回 false 时停止，表示提前停止迭代。

Go 标准库 iter 包 中定义了 Seq 和 Seq2 作为迭代器的简称，它们将每个序列元素的 1 或 2 个值传递给 yield 函数：

Go
1
2
3
4
5
6
// Seq 是 sequence 的缩写，因为迭代器会循环遍历一系列值
// Seq2 表示成对值的序列，通常是键值对或索引值对
type (
    Seq[V any]     func(yield func(V) bool)
    Seq2[K, V any] func(yield func(K, V) bool)
)

自己实现一个迭代器

实现一个容器

这里我们以乐队容器为例，演示如何为自己实现的容器，写一个迭代器。

Go
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
// 支持任意索引的乐队列表，比如将乐队 id 或 Name 作为索引
type BandList[E comparable] struct {
    band map[E]Band
}

type Band struct {
    ID     uint
    Name   string
    Since  uint
    Member map[string]string
    Songs  []string
}

func (b *BandList[E]) Add(k E, v Band) {
    b.band[k] = v
}

func NewBandList[E comparable]() *BandList[E] {
    return &BandList[E]{band: make(map[E]Band)}
}

Push 迭代器

Push 迭代器将序列中的每个值推送到 yield 函数。Push 迭代器是 Go 标准库中的标准迭代器，并由 for/range 语句直接支持。

为先前定义的容器类型实现一个返回所有元素的迭代器：

Go
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
// 迭代器函数或者方法通常根据被遍历的序列命名，例如用 All 方法返回序列中所有的元素。

// 返回 Seq 「键迭代器」
func (s *BandList[E]) Keys() iter.Seq[E] {
    return func(yield func(E) bool) {
        for v := range s.band {
            if !yield(v) {
                return
            }
        }
    }
}
// 返回 Seq2 「键值对迭代器」
func (s *BandList[E]) All() iter.Seq2[E, Band] {
    return func(yield func(E, Band) bool) {
        for k, v := range s.band {
            if !yield(k, v) {
                return
            }
        }
    }
}

使用 for/range 进行遍历：

Go
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
func main() {
    bandlist := NewBandList[string]()
    band1 := Band{
        ID:    1,
        Name:  "DOUODU",
        Since: 2022,
        Member: map[string]string{
            "杜豆豆": "主唱/键盘",
            "杜咪咪": "电子/打击乐",
        },
        Songs: []string{"南三环东路", "当我们点起火把"},
    }
    band2 := Band{
        ID:    2,
        Name:  "福禄寿",
        Since: 2022,
        Member: map[string]string{
            "杜豆豆": "主唱/键盘",
            "杜咪咪": "电子/打击乐",
        },
        Songs: []string{"我用什么把你留住", "马"},
    }
    bandlist.Add(band1.Name, band1)
    bandlist.Add(band2.Name, band2)

    // iter.Seq Push 迭代器遍历 bandlist
    for k := range bandlist.Keys() {
        fmt.Println(k)
    }

    // iter.Seq2 Push 迭代器遍历 bandlist
    for k, v := range bandlist.All() {
        fmt.Println(k, v)
    }

    // 对于一些简单的场景，比如只是想打印下元素
    // 可以不使用 for/range 语句遍历迭代器
    // 直接把函数传入迭代器，相当于自己实现了一个 yield 函数
    bandlist.Keys()(func(v string) bool {
        fmt.Println(v)
        return true
    })
}

Pull 迭代器

Pull 迭代器的工作方式则相反。每次调用 Pull 迭代器时，它都会从序列中拉出另一个值并返回该值。for/range 语句不直接支持 Pull 迭代器；但可以通过编写一个简单的 for 循环遍历 Pull 迭代器。

通常不需要自行实现一个 Pull 迭代器，新的标准库中 iter.Pull 和 iter.Pull2 函数能够将标准 Push 迭代器转为 Pull 迭代器。

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

Pull/Pull2 会返回两个函数：

  • next()：每次调用都会返回序列中的下一个值和一个布尔值，该布尔值表示该值是否有效；
  • stop()：应在完成 Pull 迭代器后调用。
Go
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
// 将标准 Push 迭代器转换为 Pull 迭代器
next, stop := iter.Pull2(bandlist.All())
defer stop()

//新建一个 bandlist，将原来的 bandlist 做一个克隆，并且把 Since 改为 2018
bandlistCopy := NewBandList[string]()
for {
    _, v, ok := next()
    if !ok {
        break
    }
    v.Since = 2018
    bandlistCopy.Add(v.Name, v)
}

// iter.Seq2 Push 迭代器遍历 bandlist
for k, v := range bandlistCopy.All() {
    fmt.Println(k, v)
}

适配器

迭代器的标准定义的一个优点是能够编写使用它们的标准适配器函数。

例如，可以定义一个过滤值序列并返回新序列的 Filter 函数。

这个 Filter 函数接受迭代器作为参数，并返回一个新的迭代器。它另一个参数是一个过滤器函数，它决定过滤器返回的新迭代器中应该过滤掉哪些值。

Go
1
2
3
4
5
6
7
8
9
func Filter[K, V any](f func(V) bool, s iter.Seq2[K, V]) iter.Seq2[K, V] {
    return func(yield func(K, V) bool) {
        for k, v := range s {
            if f(v) && !yield(k, v) {
                return
            }
        }
    }
}
Go
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
// 过滤 2017 年之后成立的乐队
for k, v := range Filter(func(b Band) bool {
    if b.Since > 2017 {
        return true
    } else {
        return false
    }
}, bandlist.All()) {
    fmt.Println(k, v)
}

总结

虽然看起来 Go 里面的迭代器很复杂（有太多的嵌套），但是只要理解了 yield 函数的作用，其原理还是比较简单的。

至于执行效率的问题，Go 编译器和运行时中做了相当多的工作来提高效率，并正确处理循环中的 break 或 panic。

评论