Go 博客

函数类型的 Range 循环

Ian Lance Taylor
2024 年 8 月 20 日

引言

这是我在 GopherCon 2024 演讲的博客文章版本。

函数类型的 Range 循环是 Go 1.23 版本中的一项新语言特性。这篇博客文章将解释我们为什么要添加这项新特性,它到底是什么,以及如何使用它。

为什么?

自 Go 1.18 以来,我们已经能够在 Go 中编写新的泛型容器类型。例如,我们来考虑这个非常简单的 `Set` 类型,它是一个基于映射实现的泛型类型。

// Set holds a set of elements.
type Set[E comparable] struct {
    m map[E]struct{}
}

// New returns a new [Set].
func New[E comparable]() *Set[E] {
    return &Set[E]{m: make(map[E]struct{})}
}

一个 Set 类型自然有添加元素和检查元素是否存在的方法。这里的细节并不重要。

// Add adds an element to a set.
func (s *Set[E]) Add(v E) {
    s.m[v] = struct{}{}
}

// Contains reports whether an element is in a set.
func (s *Set[E]) Contains(v E) bool {
    _, ok := s.m[v]
    return ok
}

除此之外,我们还需要一个函数来返回两个 Set 的并集。

// Union returns the union of two sets.
func Union[E comparable](s1, s2 *Set[E]) *Set[E] {
    r := New[E]()
    // Note for/range over internal Set field m.
    // We are looping over the maps in s1 and s2.
    for v := range s1.m {
        r.Add(v)
    }
    for v := range s2.m {
        r.Add(v)
    }
    return r
}

我们来看一下这个 `Union` 函数的实现。为了计算两个 Set 的并集,我们需要一种方法来获取每个 Set 中的所有元素。在这段代码中,我们使用 for/range 语句遍历 Set 类型的一个未导出字段。这只有在 `Union` 函数定义在 Set 包中时才有效。

但是,人们可能有很多理由需要遍历 Set 中的所有元素。这个 Set 包必须为其用户提供一种方法来做到这一点。

那该如何实现呢?

推送 Set 元素

一种方法是提供一个 `Set` 方法,该方法接受一个函数,并用 Set 中的每个元素调用该函数。我们称之为 `Push`,因为 `Set` 将每个值推送到该函数。如果该函数返回 false,我们就停止调用它。

func (s *Set[E]) Push(f func(E) bool) {
    for v := range s.m {
        if !f(v) {
            return
        }
    }
}

在 Go 标准库中,我们看到这种通用模式用于 `sync.Map.Range` 方法、`flag.Visit` 函数和 `filepath.Walk` 函数等情况。这是一种通用模式,而不是精确模式;碰巧的是,这三个示例的工作方式都不完全相同。

这是使用 `Push` 方法打印 Set 中所有元素的样子:你用一个函数调用 `Push`,该函数对元素执行你想要的操作。

func PrintAllElementsPush[E comparable](s *Set[E]) {
    s.Push(func(v E) bool {
        fmt.Println(v)
        return true
    })
}

拉取 Set 元素

另一种遍历 `Set` 元素的方法是返回一个函数。每次调用该函数时,它都会从 `Set` 返回一个值,以及一个报告该值是否有效的布尔值。当循环遍历所有元素后,布尔结果将为 false。在这种情况下,我们还需要一个在不再需要值时可以调用的停止函数。

此实现使用一对通道,一个用于 Set 中的值,另一个用于停止返回值。我们使用 goroutine 在通道上发送值。`next` 函数通过从元素通道读取来返回 Set 中的一个元素,而 `stop` 函数通过关闭停止通道来告诉 goroutine 退出。我们需要 `stop` 函数来确保在不再需要值时 goroutine 退出。

// Pull returns a next function that returns each
// element of s with a bool for whether the value
// is valid. The stop function should be called
// when finished calling the next function.
func (s *Set[E]) Pull() (func() (E, bool), func()) {
    ch := make(chan E)
    stopCh := make(chan bool)

    go func() {
        defer close(ch)
        for v := range s.m {
            select {
            case ch <- v:
            case <-stopCh:
                return
            }
        }
    }()

    next := func() (E, bool) {
        v, ok := <-ch
        return v, ok
    }

    stop := func() {
        close(stopCh)
    }

    return next, stop
}

标准库中没有任何东西完全以这种方式工作。`runtime.CallersFrames` 和 `reflect.Value.MapRange` 都类似,尽管它们返回带有方法的值,而不是直接返回函数。

这是使用 `Pull` 方法打印 `Set` 中所有元素的样子。你调用 `Pull` 来获取一个函数,然后你在 for 循环中重复调用该函数。

func PrintAllElementsPull[E comparable](s *Set[E]) {
    next, stop := s.Pull()
    defer stop()
    for v, ok := next(); ok; v, ok = next() {
        fmt.Println(v)
    }
}

标准化方法

我们现在已经看到了两种遍历 Set 中所有元素的不同方法。不同的 Go 包使用这些方法和其他几种方法。这意味着当你开始使用一个新的 Go 容器包时,你可能需要学习一种新的循环机制。这也意味着我们无法编写一个适用于几种不同类型容器的函数,因为容器类型将以不同的方式处理循环。

我们希望通过开发遍历容器的标准方法来改进 Go 生态系统。

迭代器

当然,这是许多编程语言中都会出现的问题。

流行于 1994 年首次出版的 《设计模式》一书将此描述为迭代器模式。你使用迭代器“提供一种按顺序访问聚合对象的元素而不暴露其底层表示的方法。”引文中所谓的聚合对象就是我一直称之为容器的东西。聚合对象或容器只是一个持有其他值的实体,就像我们一直在讨论的 `Set` 类型一样。

和编程中的许多想法一样,迭代器可以追溯到 Barbara Liskov 在 20 世纪 70 年代开发的 CLU 语言

如今,许多流行语言都以某种方式提供了迭代器,其中包括 C++、Java、Javascript、Python 和 Rust。

然而,Go 在 1.23 版本之前没有。

For/range

众所周知,Go 拥有内置的容器类型:切片、数组和映射。它还有一种不暴露底层表示即可访问这些值元素的方法:for/range 语句。for/range 语句适用于 Go 的内置容器类型(以及字符串、通道,从 Go 1.22 开始,还有 int)。

for/range 语句是迭代,但它不是当今流行语言中出现的迭代器。尽管如此,如果能够使用 for/range 迭代用户定义的容器(如 `Set` 类型),那将是非常好的。

然而,Go 在 1.23 版本之前不支持此功能。

此版本中的改进

对于 Go 1.23,我们决定支持对用户定义容器类型的 for/range 循环以及迭代器的标准化形式。

我们扩展了 for/range 语句以支持对函数类型的循环。我们将在下面看到这如何有助于遍历用户定义的容器。

我们还添加了标准库类型和函数来支持将函数类型用作迭代器。迭代器的标准定义使我们能够编写与不同容器类型平稳协作的函数。

函数类型的 Range 循环(部分)

改进后的 for/range 语句不支持任意函数类型。从 Go 1.23 开始,它现在支持对接受单个参数的函数进行循环。该单个参数本身必须是一个接受零到两个参数并返回布尔值的函数;按照惯例,我们称之为 yield 函数。

func(yield func() bool)

func(yield func(V) bool)

func(yield func(K, V) bool)

当我们在 Go 中谈论迭代器时,我们指的是具有这三种类型之一的函数。正如我们将在下面讨论的,标准库中还有另一种迭代器:拉取迭代器。当需要区分标准迭代器和拉取迭代器时,我们将标准迭代器称为推送迭代器。这是因为,正如我们将看到的,它们通过调用 yield 函数来推出一系列值。

标准(推送)迭代器

为了使迭代器更易于使用,新的标准库包 `iter` 定义了两种类型:`Seq` 和 `Seq2`。这些是迭代器函数类型的名称,即可用于 for/range 语句的类型。名称 `Seq` 是 sequence 的缩写,因为迭代器循环遍历一系列值。

package iter

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

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

// for now, no Seq0

`Seq` 和 `Seq2` 之间的区别仅在于 `Seq2` 是一个对序列,例如映射中的键和值。在这篇文章中,为简化起见,我们将重点关注 `Seq`,但我们所说的大部分内容也适用于 `Seq2`。

用例子解释迭代器的工作原理最简单。这里 `Set` 方法 `All` 返回一个函数。`All` 的返回类型是 `iter.Seq[E]`,所以我们知道它返回一个迭代器。

// All is an iterator over the elements of s.
func (s *Set[E]) All() iter.Seq[E] {
    return func(yield func(E) bool) {
        for v := range s.m {
            if !yield(v) {
                return
            }
        }
    }
}

迭代器函数本身接受另一个函数(yield 函数)作为参数。迭代器用 Set 中的每个值调用 yield 函数。在这种情况下,迭代器,即 `Set.All` 返回的函数,与我们之前看到的 `Set.Push` 函数非常相似。

这展示了迭代器的工作方式:对于某些值序列,它们用序列中的每个值调用 yield 函数。如果 yield 函数返回 false,则不再需要值,迭代器可以返回,并执行任何可能需要的清理。如果 yield 函数从不返回 false,迭代器可以在用序列中的所有值调用 yield 后返回。

这就是它们的工作方式,但我们承认,当你第一次看到它们时,你的第一反应可能是“这里有很多函数在飞来飞去”。你没有说错。我们关注两件事。

首先,一旦你越过此函数代码的第一行,迭代器的实际实现就非常简单:用 Set 的每个元素调用 yield,如果 yield 返回 false 则停止。

        for v := range s.m {
            if !yield(v) {
                return
            }
        }

其次,使用它真的很容易。你调用 `s.All` 获取一个迭代器,然后使用 for/range 遍历 `s` 中的所有元素。for/range 语句支持任何迭代器,这显示了它有多么容易使用。

func PrintAllElements[E comparable](s *Set[E]) {
    for v := range s.All() {
        fmt.Println(v)
    }
}

在这种代码中,`s.All` 是一个返回函数的的方法。我们调用 `s.All`,然后使用 for/range 遍历它返回的函数。在这种情况下,我们可以让 `Set.All` 本身就是一个迭代器函数,而不是让它返回一个迭代器函数。然而,在某些情况下这不起作用,例如如果返回迭代器的函数需要接受一个参数,或者需要做一些设置工作。按照惯例,我们鼓励所有容器类型提供一个返回迭代器的 `All` 方法,这样程序员就不必记住是直接遍历 `All` 还是调用 `All` 来获取可以遍历的值。他们总是可以做后者。

如果你仔细想想,你会发现编译器必须调整循环以创建一个 yield 函数来传递给 `s.All` 返回的迭代器。Go 编译器和运行时中存在相当多的复杂性,以使其高效运行,并正确处理循环中的 `break` 或 `panic` 等情况。我们不会在这篇博客文章中讨论这些。幸运的是,在实际使用此功能时,实现细节并不重要。

拉取迭代器

我们现在已经了解了如何在 for/range 循环中使用迭代器。但简单的循环并不是使用迭代器的唯一方式。例如,有时我们可能需要并行遍历两个容器。我们该如何做到这一点呢?

答案是使用另一种迭代器:拉取迭代器。我们已经看到,标准迭代器(也称为推送迭代器)是一个接受 yield 函数作为参数并通过调用 yield 函数推送序列中每个值的函数。

拉取迭代器的工作方式则相反:它是一个每次调用它都会返回序列中下一个值的函数。

我们将重复两种迭代器之间的区别以帮助您记忆

  • 推送迭代器将序列中的每个值推送到一个 yield 函数。推送迭代器是 Go 标准库中的标准迭代器,并直接由 for/range 语句支持。
  • 拉取迭代器的工作方式则相反。每次调用拉取迭代器时,它都会从序列中拉取另一个值并返回。拉取迭代器不直接受 for/range 语句支持;但是,编写一个遍历拉取迭代器的普通 for 语句是直截了当的。实际上,我们之前在查看使用 `Set.Pull` 方法时看到了一个例子。

你可以自己编写一个拉取迭代器,但通常不必这样做。新的标准库函数 `iter.Pull` 接受一个标准迭代器,也就是说一个推送迭代器函数,并返回一对函数。第一个是拉取迭代器:一个每次调用都会返回序列中下一个值的函数。第二个是停止函数,当我们完成拉取迭代器时应该调用它。这与我们之前看到的 `Set.Pull` 方法类似。

`iter.Pull` 返回的第一个函数,即拉取迭代器,返回一个值和一个布尔值,用于报告该值是否有效。在序列的末尾,布尔值将为 false。

`iter.Pull` 返回一个停止函数,以防我们没有遍历到序列的末尾。在一般情况下,推送迭代器(`iter.Pull` 的参数)可能会启动 goroutine,或者构建需要在迭代完成后清理的新数据结构。当 yield 函数返回 false(表示不再需要值)时,推送迭代器将执行任何清理。当与 for/range 语句一起使用时,for/range 语句将确保如果循环提前退出(通过 `break` 语句或任何其他原因),则 yield 函数将返回 false。另一方面,对于拉取迭代器,无法强制 yield 函数返回 false,因此需要停止函数。

另一种说法是,调用停止函数将导致 yield 函数在被推送迭代器调用时返回 false。

严格来说,如果拉取迭代器返回 false 以指示它已到达序列末尾,则无需调用停止函数,但通常始终调用它更简单。

这是一个使用拉取迭代器并行遍历两个序列的示例。此函数报告两个任意序列是否包含相同顺序的相同元素。

// EqSeq reports whether two iterators contain the same
// elements in the same order.
func EqSeq[E comparable](s1, s2 iter.Seq[E]) bool {
    next1, stop1 := iter.Pull(s1)
    defer stop1()
    next2, stop2 := iter.Pull(s2)
    defer stop2()
    for {
        v1, ok1 := next1()
        v2, ok2 := next2()
        if !ok1 {
            return !ok2
        }
        if ok1 != ok2 || v1 != v2 {
            return false
        }
    }
}

该函数使用 `iter.Pull` 将两个推送迭代器 `s1` 和 `s2` 转换为拉取迭代器。它使用 `defer` 语句确保在我们完成拉取迭代器时它们被停止。

然后代码循环,调用拉取迭代器检索值。如果第一个序列已完成,则如果第二个序列也已完成,则返回 true,否则返回 false。如果值不同,则返回 false。然后它循环以拉取接下来的两个值。

与推送迭代器一样,Go 运行时在使拉取迭代器高效方面存在一些复杂性,但这不影响实际使用 `iter.Pull` 函数的代码。

迭代器上的迭代

现在你已经了解了关于函数类型的 range 循环和迭代器的所有知识。我们希望你喜欢使用它们!

然而,还有一些事情值得一提。

适配器

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

例如,这是一个过滤值序列并返回新序列的函数。这个 `Filter` 函数接受一个迭代器作为参数并返回一个新的迭代器。另一个参数是一个过滤函数,它决定哪些值应该包含在 `Filter` 返回的新迭代器中。

// Filter returns a sequence that contains the elements
// of s for which f returns true.
func Filter[V any](f func(V) bool, s iter.Seq[V]) iter.Seq[V] {
    return func(yield func(V) bool) {
        for v := range s {
            if f(v) {
                if !yield(v) {
                    return
                }
            }
        }
    }
}

与之前的示例一样,当你第一次看到函数签名时,它们看起来很复杂。一旦你越过签名,实现就非常简单了。

        for v := range s {
            if f(v) {
                if !yield(v) {
                    return
                }
            }
        }

代码遍历输入迭代器,检查过滤函数,并用应该进入输出迭代器的值调用 yield。

我们将在下面展示一个使用 `Filter` 的示例。

(Go 标准库中目前没有 `Filter` 版本,但未来版本可能会添加。)

二叉树

作为一个例子,说明推送迭代器在遍历容器类型时有多么方便,让我们考虑这个简单的二叉树类型。

// Tree is a binary tree.
type Tree[E any] struct {
    val         E
    left, right *Tree[E]
}

我们不会展示将值插入树中的代码,但自然地应该有一种方法来遍历树中的所有值。

事实证明,如果迭代器代码返回一个布尔值,编写起来会更容易。由于 for/range 支持的函数类型不返回任何东西,所以这里的 `All` 方法返回一个小的函数字面量,它调用迭代器本身(这里称为 `push`),并忽略布尔结果。

// All returns an iterator over the values in t.
func (t *Tree[E]) All() iter.Seq[E] {
    return func(yield func(E) bool) {
        t.push(yield)
    }
}

// push pushes all elements to the yield function.
func (t *Tree[E]) push(yield func(E) bool) bool {
    if t == nil {
        return true
    }
    return t.left.push(yield) &&
        yield(t.val) &&
        t.right.push(yield)
}

`push` 方法使用递归遍历整个树,对每个元素调用 yield。如果 yield 函数返回 false,则该方法一路返回 false。否则,一旦迭代完成,它就返回。

这展示了使用这种迭代器方法遍历复杂数据结构是多么简单。无需维护单独的堆栈来记录树中的位置;我们可以直接使用 goroutine 调用堆栈来完成此操作。

新的迭代器函数。

Go 1.23 中新增的还有切片和映射包中与迭代器一起使用的函数。

以下是切片包中的新函数。`All` 和 `Values` 是返回切片元素迭代器的函数。`Collect` 从迭代器中获取值并返回一个包含这些值的切片。其他函数请参阅文档。

以下是 maps 包中的新函数。`All`、`Keys` 和 `Values` 返回映射内容的迭代器。`Collect` 从迭代器中获取键和值并返回一个新的映射。

标准库迭代器示例

这里是一个如何将这些新函数与我们之前看到的 `Filter` 函数一起使用的示例。此函数接受一个 int 到 string 的映射,并返回一个切片,该切片仅包含映射中长度大于参数 `n` 的值。

// LongStrings returns a slice of just the values
// in m whose length is n or more.
func LongStrings(m map[int]string, n int) []string {
    isLong := func(s string) bool {
        return len(s) >= n
    }
    return slices.Collect(Filter(isLong, maps.Values(m)))
}

`maps.Values` 函数返回 `m` 中值的迭代器。`Filter` 读取该迭代器并返回一个只包含长字符串的新迭代器。`slices.Collect` 从该迭代器读取到新切片中。

当然,你可以很容易地编写一个循环来完成此操作,并且在许多情况下,循环会更清晰。我们不希望鼓励每个人总是以这种风格编写代码。也就是说,使用迭代器的好处是这种函数以相同的方式与任何序列一起工作。在这个例子中,请注意 Filter 如何使用映射作为输入和切片作为输出,而无需更改 Filter 中的代码。

遍历文件中的行

尽管我们看到的大多数示例都涉及容器,但迭代器是灵活的。

考虑这段简单的代码,它不使用迭代器,用于遍历字节切片中的行。这很容易编写,并且效率相当高。

    nl := []byte{'\n'}
    // Trim a trailing newline to avoid a final empty blank line.
    for _, line := range bytes.Split(bytes.TrimSuffix(data, nl), nl) {
        handleLine(line)
    }

然而,`bytes.Split` 确实分配并返回一个字节切片切片来存储行。垃圾回收器将不得不做一些工作才能最终释放该切片。

这是一个返回字节切片行迭代器的函数。除了通常的迭代器签名,这个函数相当简单。我们不断从数据中选择行,直到没有剩余,然后我们将每一行传递给 yield 函数。

// Lines returns an iterator over lines in data.
func Lines(data []byte) iter.Seq[[]byte] {
    return func(yield func([]byte) bool) {
        for len(data) > 0 {
            line, rest, _ := bytes.Cut(data, []byte{'\n'})
            if !yield(line) {
                return
            }
            data = rest
        }
    }
}

现在我们遍历字节切片行的代码看起来像这样。

    for line := range Lines(data) {
        handleLine(line)
    }

这和之前的代码一样容易编写,而且效率更高,因为它不需要分配一个行切片。

将函数传递给推送迭代器

对于我们的最后一个例子,我们将看到你不需要在 range 语句中使用推送迭代器。

之前我们看到了一个 `PrintAllElements` 函数,它打印出集合的每个元素。这是打印集合所有元素的另一种方法:调用 `s.All` 获取一个迭代器,然后传入一个手写的 yield 函数。这个 yield 函数只是打印一个值并返回 true。请注意,这里有两个函数调用:我们调用 `s.All` 获取一个迭代器,该迭代器本身就是一个函数,然后我们用我们手写的 yield 函数调用该函数。

func PrintAllElements[E comparable](s *Set[E]) {
    s.All()(func(v E) bool {
        fmt.Println(v)
        return true
    })
}

没有特别的理由以这种方式编写此代码。这只是一个例子,表明 yield 函数并非神奇。它可以是任何你喜欢的函数。

更新 go.mod

最后一点:每个 Go 模块都指定了它使用的语言版本。这意味着为了在现有模块中使用新的语言特性,你可能需要更新该版本。这适用于所有新的语言特性;它不是函数类型循环特有的。由于 Go 1.23 版本中新增了函数类型循环,因此使用它需要至少指定 Go 语言版本 1.23。

有(至少)四种方法可以设置语言版本

  • 在命令行上,运行 `go get go@1.23`(或 `go mod edit -go=1.23` 仅编辑 `go` 指令)。
  • 手动编辑 `go.mod` 文件并更改 `go` 行。
  • 保留模块整体的旧语言版本,但使用 `//go:build go1.23` 构建标签允许在特定文件中使用函数类型循环。

下一篇文章:新的 unique 包
上一篇文章:Go 1.23 已发布
博客索引