基于 Go 泛型实现迭代器

作者 | Kulics
整理 | Hana

本文作者是知乎@Kulics，经授权后由编程语言 Lab 转载发布。

Kulics，Feel 语言、XyKey 应用作者。

代码实现：https://github.com/kulics/Gollection

迭代器模式

迭代器模式（Iterator Pattern） 是一种非常常用的设计模式。这种模式在 Java、C#、Rust 等语言中都有大量应用，这种模式可以以一种相同的方式，顺序访问集合类型中的元素，而无需关注集合类型实现方式。

应用迭代器模式，我们就可以对无论是数组、链表或是树形式实现的集合类型统一遍历，甚至组合更多高级函数功能使用，提升代码的表达能力。

Go 在之前没有泛型的时候，很难以类型安全的方式抽象迭代器模式出来，只能靠内置的 for range 来对 slice、map 类型遍历，我们很难扩展其它的自定义类型支持遍历。

现在有了泛型也就有了这个能力来让我们自己实现通用的遍历方式。

下面来让我们看看如何在 Go 中实现这样的迭代器。

Iterator

关于 Go 的泛型语法不是本篇文章的重点，就不再赘述，我们直接从 Iterator 的定义开始。

要让一个类型支持遍历很简单，熟悉 Go 的同学们可能过去或多或少都用过类似的技巧。

让我们先从简单的 ForEach 函数开始，假设我们有一个使用 slice 实现的 List 类型，那么我们只需要利用 Go 的函数一等公民特性，就能让外部来遍历这个类型。

我们让 ForEach 传入一个函数类型的参数，来让调用者可以通过外部函数的方式来遍历 List 里面的元素。

type MyList[T any] struct {
  elements []T
}

func (a MyList[T]) ForEach(f func(T)) {
  for _, v := range a.elements {
    f(v) // 交由调用者来定义遍历的动作
  }
}

func main() {
  var list = MyList[int]{[]int{1, 2, 3, 4, 5}}
  list.ForEach(func(i int) {
    println(i)
  })
}

这个方式很简单，速度通常也很快，如果我们的集合类型都可以提供这样一个函数，也能支持一样的遍历。

所以我们很容易就能得出这样一种迭代器模型：

type Iterator[T any] interface {
  ForEach(f func(T))
}

这样任何一个类型只要提供了 ForEach，我们就可以用来遍历。

但这样就够了吗？emmmmm，还稍微差了点。

仔细想想我们就能发现，这种 ForEach 只能让所有集合类型立即执行遍历行为，如果我们想组合一些处理，就会带来很多额外的遍历，对性能造成负影响。

如果迭代器可以每次调用都只返回下一个元素，直到无元素时停止，对我们组合各种操作就会便利得多了。

我们可以梳理出这样两个函数，一个用来告诉我们还有没有下一个元素，另一个用来取得下一个元素，这样就可以组成迭代器的两个基本操作了，事实上 java 的迭代器也是这样设计的。

type Iterator[T any] interface {
  HasNext() bool
  Next() T
}

这样我们拿到一个迭代器之后，使用简单的 for 循环就能完成遍历了，比 ForEach 更方便的是，现在我们可以随时中断了。这里先假设 MyList 已经实现了 HasNext 和 Next。

func main() {
  var iter Iterator[int] = MyList[int]{[]int{1, 2, 3, 4, 5}}
  for iter.HasNext() { // 判断还有下一个
    var item = iter.Next()
    if item == 3 {
      break // 提前中断
    }
    println(item)
  }
}

当然对于 Go 来说，我们可以在一次调用的时候返回多个值，所以其实可以在 Next 的同时就判断是不是还有下一个元素，所以我们可以将 HasNext 和 Next 合并成一个 Next 函数，相信看了下面的代码后同学们都会觉得很熟悉。

type Iterator[T any] interface {
  Next() (v T, ok bool)
}

func main() {
  var iter Iterator[int] = MyList[int]{[]int{1, 2, 3, 4, 5}}
  for item, ok := iter.Next(); ok; item, ok = iter.Next() { // 效果是等价的
    println(item)
  }
}

Iterable

看到这里相信大家大概对迭代器有个基本的概念了，可能有部分同学已经发现了 Next 的问题了。

是的，Next 是单向的，一旦开始调用，就只能一直向后移动直到终结了。

这样的 Iterator 还不适合作为集合类型的通用接口，因为我们当然不希望 MyList 构造出来只能遍历一次了。

所以我们需要另一个函数，来让我们可以反复生成不同的独立的 Iterator，这样就可以反复遍历了。

这个函数很简单，就只要给我们取得 Iterator 就可以了，这里我们借鉴 java 的风格称之为 Iterable。

type Iterable[T any] interface {
  Iter() Iterator[T]
}

通过 Iterable 和 Iterator 的类型搭配，我们就可以实现出一种通用的方式来实现集合类型遍历的方法。

所有 List、Map、Set 等数据结构都可以应用这种模式，下面直接给出 MyList 的简要实现。

type MyList[T any] struct {
  elements []T
}

func (a MyList[T]) Iter() Iterator[T] {
  return &MyListIter[T]{a, -1}
}

type Iterator[T any] interface {
  Next() (v T, ok bool)
}

type MyListIter[T any] struct {
  list MyList[T]
  index int
}

func (a *MyListIter[T]) Next() (v T, ok bool) {
  var null T
  if a.index <len(a.list.elements)-1 {
    a.index++
    return a.list.elements[a.index], true
  }
  return null, false
}

组合操作

我们有了迭代器，用 for 来遍历已经没问题了，但我们前面还提到过的组合操作又是什么回事呢？

别急，我们现在来看看。

ForEach 对于我们来说非常熟悉了，我们现在可以基于迭代器来实现更通用的 ForEach 方法。

对于任意的 Iterator，我们就用直接遍历就实现了。

func ForEach[T any](f func(T), iter Iterator[T]) {
  for item, ok := iter.Next(); ok; item, ok = iter.Next() {
    f(item)
  }
}

func main() {
  var list = MyList[int]{[]int{1, 2, 3, 4, 5}}
  ForEach(func(i int) {
    println(i)
  }, list.Iter())
}

是不是很方便？在 Go1.18 给我们带来了泛型之后，我们终于不需要用骚操作来实现一个通用的 ForEach 了。

ForEach 是一种典型的终结操作，终结操作的意思是到这一步就已经最终执行遍历来得到我们想要的结果了，这些函数通常是某一段操作的最后一环，同类型的还有 Contains、Sum、Average、Max、Min 之类的，大家有兴趣可以自己尝试实现一下。

相比于终结操作，还有另一类操作是转换操作。很显然，这类操作就是那类不会立刻执行遍历，而是增加一些中间处理的函数，比如从一个迭代器映射到另一个迭代器，通常称之为 Map；或者是从一个迭代器中筛选出我们需要的元素，通常称之为 Filter；同类的还有 Step、Concat、Index 等。

我们来看看最典型的 Map 是如何实现的。

很简单，既然我们每次调用 Next 就可以取得下一个元素，那么我们只要将这个元素传给处理函数，就能得到一个映射后的另一个值。

我们先用伪代码看看转换函数是怎么工作的：

var iter Iterator[T] = ......
var transform func(T) R = ...... // 假设我们有一个将 T 类型转换到 R 类型的函数
item, ok := iter.Next()
if ok {
  var newItem R = transform(item) // 现在我们得到了一个映射后的值
}

既然我们只需要对 Next 函数调用一下 transform，那我们只需要一个将 transform 存到一个新的迭代器里，然后这个新迭代器每次 Next 时调用 transform 就可以实现完美一对一的映射了。

下面我们看看一个最简单的实现：

type MapIter[T any, R any] struct {
  transform func(T) R
    iterator  Iterator[T]
}

func (a *MapIter[T, R]) Next() (R, bool) { // 注意新的 MapIter 实现的 Next 返回的是 R 而不是 T
  var null R
  if v, ok := a.iterator.Next(); ok {
    return a.transform(v), true
  }
  return null, false
}

有了 MapIter，我们包装一下 Map 函数就能用了。

func Map[T any, R any](transform func(T) R, iter Iterator[T]) Iterator[R] {
  return &MapIter[T, R]{transform, iter}
}

非常简单易用，我们看看使用效果吧。

func main() {
  var list = MyList[int]{[]int{1, 2, 3, 4, 5}}
  ForEach(func(i bool) {
    println(i)
  }, Map(func(i int) bool {
    return i%2 == 0
  }, list.Iter()))
}