Go1.23新特性: 迭代器-在函数上遍历

 官方文档： Range Over Function Types - The Go Programming Language

go1.23发布一个新特性，range语句，除了接收内置的容器之外，支持接收一个迭代器函数来实现对元素的遍历，函数的可以是形式如下三种，分别表示每次迭代获取0/1/2个元素。

func(func() bool)
func(func(K) bool)
func(func(K, V) bool)

1. 为什么需要

从go1.18开始，我们可以使用泛型(generic)来实现支持不同元素类型的容器，考虑一个集合(Set)的例子，在map的基础上实现集合的功能：

package set

type Set[E comparable] struct {
	m map[E]struct{}
}

func New[E comparable]() *Set[E] {
	return &Set[E]{
		make(map[E]struct{}),
	}
}

func (s *Set[E]) Add(v E) {
	s.m[v] = struct{}{}
}

func (s *Set[E]) Contains(v E) bool {
	_, ok := s.m[v]
	return ok
}

下面我们实现一个获取两个集合并集的功能：

func Union[E comparable](s1, s2 *Set[E]) *Set[E] {
	s := New[E]()
	for v := range s1.m {
		s.Add(v)
	}
	for v := range s2.m {
		s.Add(v)
	}
	return s
}

这里为了实现并集的功能，需要分别遍历两个集合的元素。如果直接遍历内部未导出的成员，需要并集功能的代码也在定义这个集合的包里。

但是我们难免会遇到在包外要遍历集合所有元素的场景，那么怎么实现呢？

1. Push模式

Set提供一个方法，接收一个作用在集合成员类型的函数，在集合的每个元素上执行这个函数。另外还可以要求这个函数返回一个布尔值，当函数返回false时，终止遍历。

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

这种模式的实现和使用示例如下。

func TestPrint(t *testing.T) {
	s := New[int]()
	s.Add(1)
	s.Add(3)
	s.Add(2)
	s.Add(4)
	s.Push(func(i int) bool {
		t.Log(i)
		return true
	})
}

一些go的标准库使用了这种模式，如sync.Map.Range、flag.Visit、filepath.Walk，当然，它们的使用细节会有不同。

这种方法实现并集的示例代码如下：

func UnionByPush[E comparable](s1, s2 *Set[E]) *Set[E] {
	s := New[E]()
	s1.Push(func(e E) bool { s.Add(e); return true })
	s2.Push(func(e E) bool { s.Add(e); return true })
	return s
}

2. Pull模式

另一个方法，是由Set返回一个方法，每次调用该方法，返回一个集合里的元素，同时带一个布尔值，来标识这个元素是否合法，是否是需要遍历的，比如已经遍历完该集合所有元素后，会把这个布尔值置为false，这表示遍历结束，没有可用的元素了。另外我们可能还需要一个stop函数，调用这个函数告诉容器，不需要再遍历了。

下面的代码实现了该功能，我们使用两个channel，一个传递集合里的元素，另一个传递结束的标志。使用goroutine发送值到channel，next函数从channel读取元素并返回，stop函数关闭stop channel来通知goroutine退出。我们需要使用stop函数来保证当没有元素时，goroutine退出。

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使用了类似的方式，只是它们返回了包含方法的对象而不是返回函数。

这种模式的使用示例代码

func TestPrintPull(t *testing.T) {
	s := New[int]()
	s.Add(1)
	s.Add(3)
	s.Add(2)
	s.Add(4)
	next, stop := s.Pull()
	defer stop()
	for v, ok := next(); ok; v, ok = next() {
		t.Log(v)
	}
}

2. 标准化

现在我们已经看到了两种不同的遍历集合中所有元素的方法。不同的包可能会使用不同的方法。这会增加使用不同包的学习负担。这也意味着我们不能编写一个函数来处理几种不同类型的容器。

我们希望通过开发容器循环的标准方法来改善Go生态系统。

1. 迭代器

迭代器是许多语言都包含的功能，如C++、Java、JavaScript、Python、Rust等。

Go1.23也开始支持。

for/range语句

for range语句实现了便利内置容器(slice、array、map、channel、正整数(1.22开始))的元素的功能。

从1.23开始，for/range语句也可以应用在用户定义的容器类型上，进而推广到标准的迭代器模式。

从1.23开始，扩展了for/range语句的功能，可以应用在一些特定类型的函数上，进而实现在用户定义类型上的循环。

for/range语句支持特定的函数类型：接收一个函数作为参数，称作yield函数，该函数接收0-2个参数，并返回一个布尔值。

func(yield func() bool)

func(yield func(V) bool)

func(yield func(K, V) bool)

我们把上面类型的函数称作迭代器，我们称之为标准迭代器(standard iterator)，为了区别另一种pull迭代器，我们称标准迭代器为push迭代器，因为这些迭代器通过yield函数，推(Push)出了一系列的值。

2. 标准(Push)迭代器

为方便使用，新的标准库提供了一个iter包，里面定义了两个类型：Seq和Seq2,作为迭代器函数的别名，表示单值和双指的序列。它们可以和for.range语句配合使用。

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

使用上面Set的遍历来解释Seq的使用。

func (s *Set[E]) All() iter.Seq[E] {
	return func(yield func(E) bool) {
		for v := range s.m {
			if !yield(v) {
				break
			}
		}
	}
}

迭代器函数接收一个yield函数作为参数，迭代器在Set的每个成员上调用yield函数。当迭代结束时，yield会返回false，表示不再迭代。在这个例子里，Set.All返回一个迭代器，和前面的push函数很类似。

下面展示了迭代器的工作模式：在一些值上分别调用yield函数，当yield返回false（循环中调用break或者return），表示不需要再迭代。迭代器可以执行一些清理逻辑，然后退出；如果yield一直返回true，则可以在序列上所有的元素上调用yield，然后再退出。

func TestPushIter(t *testing.T) {
	s := New[int]()
	s.Add(1)
	s.Add(3)
	s.Add(2)
	s.Add(4)
	for v := range s.All() {
		t.Log(v)
        // yield(3) yield(4)会返回false
		if v > 2 {
			break
		}
	}
}

这个逻辑看起来有点复杂，但里面有两个要点：

1. 迭代器的实现很简单，就是在每个希望被迭代到的元素上调用yield函数，如果这个函数返回false，就退出。（如果不退出，程序会异常）。

2. 迭代器的使用也很简单，调用s.All返回一个迭代器，然后用for/range语句来循环s的元素。

在这里，Set.All方法返回了一个迭代器函数，其实我们也可以直接把Set.All自身作为一个迭代器。

不过在一些情况下会有些限制，比如需要一些参数来调整迭代器的行为、需要做一些初始化的工作等。为方便起见，我们鼓励每个容器类型提供一个All方法，返回一个迭代器函数，这样调用者不用再记忆是在All方法上迭代还是在它的返回值上迭代。

如果你仔细想想，就会发现编译器需要创建合适的yield函数，传递给s.All返回的迭代器。Go编译器和运行时中有相当多的复杂性，以便高效地循环，并正确处理循环中的break或panic等问题。我们不会在本篇博文中介绍这些内容。幸运的是，在实际使用此功能时，实现细节并不重要。

3. Pull迭代器

上面展示了如何使用for/range循环来使用迭代器。但对于一些复杂的应用，有时我们需要同时迭代两个容器，需要怎么做呢？

答案是，使用另一种迭代器: Pull迭代器。不同于前面的push迭代器接收一个yield函数，pull迭代器采用另一种方式：它也是一个函数，但你需要再每次循环里调用它，获取序列里的下一个元素。

我们再重复一下两种迭代器的区别来强化记忆：

• push迭代器把序列里的元素，推(push)到一个yield函数里，push迭代器是go的标准迭代器，被for/range直接支持。
• pull迭代器与前者不同：每次你调用迭代器，它从序列里拉(pull)一个值并返回。for/range并不支持这种迭代器，然而，编写一个普通的 for 语句来循环遍历 pull 迭代器是很简单的事情。事实上，我们在之前讨论使用 Set.Pull 方法时看到了一个例子。

您可以自己编写一个pull迭代器，但通常不必这样做。新的标准库函数iter.Pull接受一个标准迭代器，即一个推送迭代器的函数，并返回一对函数。第一个是拉式迭代器：每次调用时都会返回序列中的下一个值的函数。第二个是停止函数，应在我们完迭代时调用。这类似于我们在上面看到的 Set.Pull 方法。

iter.Pull 返回的第一个函数是 pull 迭代器，它返回一个值和一个用于报告该值是否有效的布尔值。布尔值在序列结束时将为 false。

iter.Pull 还返回一个停止函数，以便在我们不需要读完序列所有的元素时，可以提前关闭迭代器。，因为在某些时候，底层的push迭代器可能会启动 goroutine，或构建需要在迭代完成时清理的新数据结构，当yield 函数返回 false 时，推送迭代器可以执一些清理；当与 for/range 语句一起使用时，for/range 语句将确保如果循环通过 break 语句或由于任何其他原因提前退出，则yield 函数将返回 false来提醒迭代器执行清理的操作。如果使用使用iter.Pull将push迭代器包装成pull迭代器使用，则没有办法强制让yield 函数返回 false，因此需要停止函数。

另一种通知方法是，调用 stop 函数会导致在由推送迭代器调用时，yield 函数返回 false。

严格来说，如果拉取迭代器返回 false 以表明它已到达序列末尾，则无需调用 stop 函数，但通常更简单的方法是始终调用它。

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

func TestPullFromPush(t *testing.T) {
	s1 := newSet(1, 2, 3, 4)
	s2 := newSet(1, 2, 3, 4)
	next1, stop1 := iter.Pull(s1.All())
	defer stop1()
	next2, stop2 := iter.Pull(s2.All())
	defer stop2()
	if func() bool {
		for {
			v1, ok1 := next1()
			v2, ok2 := next2()
			fmt.Println(v1, v2)
			if !ok1 {
				if ok2 {
					return !ok2
				}
			}
			if ok1 != ok2 || v1 != v2 {
				return false
			}
		}
	}() {
		t.Log("True")
	} else {
		t.Log("False")
	}

}

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

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

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

迭代器上的迭代

现在您已经了解了有关函数类型范围和迭代器的所有知识。希望您喜欢使用它们！

不过，还有一些事情值得一提。

适配器

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

例如，下面的函数，对一个序列进行过滤，返回一个新序列。此 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 版本，但未来版本中可能会添加一个。）

二叉树的遍历

作为push迭代器在容器类型上进行循环的一个示例，让我们考虑一个简单的二叉树类型。

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。否则，它只会在迭代完成后返回。

这个方法实现了前序遍历，调整return行的调用顺序就可以实现中序、后序遍历。

这表明使用此迭代器方法循环遍历甚至复杂的数据结构是多么简单。无需维护单独的堆栈来记录树中的位置；我们只需使用 goroutine 调用堆栈即可完成此操作。

新的迭代器函数

Go 1.23 在slices和maps包中还新增了与迭代器配合使用的函数。

以下是slices包中的新函数。All 和 Values 是返回切片元素上的迭代器的函数。Collect 从迭代器中获取值并返回保存这些值的切片。有关其他函数，请参阅文档。

• All([]E) iter.Seq2[int, E]
• Values([]E) iter.Seq[E]
• Collect(iter.Seq[E]) []E
• AppendSeq([]E, iter.Seq[E]) []E
• Backward([]E) iter.Seq2[int, E]
• Sorted(iter.Seq[E]) []E
• SortedFunc(iter.Seq[E], func(E, E) int) []E
• SortedStableFunc(iter.Seq[E], func(E, E) int) []E
• Repeat([]E, int) []E
• Chunk([]E, int) iter.Seq([]E)

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

• All(map[K]V) iter.Seq2[K, V]
• Keys(map[K]V) iter.Seq[K]
• Values(map[K]V) iter.Seq[V]
• Collect(iter.Seq2[K, V]) map[K, V]
• Insert(map[K, V], iter.Seq2[K, V])

标准库迭代器使用示例

下面是一个示例，说明如何将这些新函数与我们之前看到的 Filter 函数一起使用。此函数接受从 int 到字符串的映射，并返回一个切片，该切片仅包含映射中长度超过某个参数 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 函数返回一个遍历map中值的迭代器。Filter 读取该迭代器并返回一个仅包含长字符串的新迭代器。slices.Collect 从该迭代器读取到新切片中。

当然，您可以编写一个循环来轻松完成此操作，而且在许多情况下，循环会更清晰。我们不想鼓励每个人一直以这种风格编写代码。也就是说，使用迭代器的优势在于这种函数对任何序列的工作方式都相同。在此示例中，请注意 Filter 如何使用 map 作为输入并使用切片作为输出，而无需更改 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)
}

这种方法使用上一样简单，但更高效一些，因为不需要为所有的行同时分配空间。

向push迭代器传一个函数参数

最后一个例子，我们将看到您不一定要在for/range语句中使用推送迭代器。

之前我们看到了一个 PrintAllElements 函数，它打印出集合的每个元素。这是打印集合所有元素的另一种方法：调用 s.All 来获取迭代器，然后传入手写的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。

有（至少）四种方法可以设置语言版本：

1. 在命令行上，运行命令

go get go@1.23

2. 运行命令

go mod edit -go=1.23

3. 手动编辑 go.mod 文件并更改 go那一行。
4. 保留整个模块的旧语言版本，但使用 //go:build go1.23 构建标签允许在特定文件中使用函数类型的遍历的功能。