Go语言rangefunc自定义迭代器的函数签名

Go 1.23 版本将 rangefunc（基于函数的 for range 循环）正式引入标准库。要编写一个能够让 for range 语句遍历的函数，必须严格遵守特定的函数签名规则。这些签名定义了迭代器如何与 Go 运行时交互，以及如何传递数据或终止循环。

1. 理解核心签名机制

rangefunc 的核心在于 for range 会调用一个特定的函数，并向其传递一个被称为 yield（产出）的回调函数。你的迭代逻辑负责调用 yield 来传递值，而 yield 的返回值决定了是否继续迭代。

掌握以下三种基础签名是编写自定义迭代器的关键：

单值迭代器签名

这是最基础的签名形式，适用于遍历只产生一个值的序列（如切片、流）。

• 定义格式：func(yield func(V) bool)
• 参数说明：
  • V：表示你要遍历的数据类型（如 int, string 或自定义结构体）。
  • yield：由 for range 传入的回调函数。
  • bool：yield 函数的返回值。如果返回 false，表示循环需要提前终止（例如使用了 break）。

键值对迭代器签名

适用于需要同时产生键和值的场景（如 Map 遍历、带索引的遍历）。

• 定义格式：func(yield func(K, V) bool)
• 参数说明：
  • K：键的类型（如 int）。
  • V：值的类型（如 string）。
  • yield：接受两个参数的回调函数。

拉取式迭代器签名 (Pull Iterator)

这是 rangefunc 的变体，通常用于 iter.Pull 或 iter.Pull2，允许手动控制迭代步进，但在直接用于 for range 时，前两者更为常见。在自定义 rangefunc 时，我们主要关注“推送”模式（Push Model），即上述两种签名。

2. 实现单值迭代器

编写一个遍历整数切片的迭代器，该函数接收 []int 并返回一个符合单值迭代器签名的函数。

1. 定义外层包装函数，接受数据源。
2. 返回一个匿名函数，该函数签名必须匹配 func(yield func(int) bool)。
3. 在返回的函数内部遍历数据。
4. 调用 yield(v) 将每个元素传给循环体。
5. 检查 yield 的返回值，若为 false 则 中断 循环。

代码示例：

// Ints 返回一个迭代器函数，用于遍历整数切片
func Ints(s []int) func(yield func(int) bool) {
    return func(yield func(int) bool) {
        for _, v := range s {
            // 调用 yield 传递值
            // 如果 yield 返回 false（例如用户在循环中执行了 break），则停止迭代
            if !yield(v) {
                return
            }
        }
    }
}

使用该迭代器：

func main() {
    data := []int{10, 20, 30}
    // 直接将函数调用放在 range 后面
    for n := range Ints(data) {
        fmt.Println(n)
    }
}

3. 实现键值对迭代器

编写一个模拟 Map 遍历或带索引遍历的迭代器。

1. 定义外层函数，决定键 K 和值 V 的类型。
2. 返回匹配 func(yield func(K, V) bool) 的函数。
3. 计算键和值。
4. 调用 yield(k, v) 并检查终止信号。

代码示例（带索引的遍历）：

// Enumerate 返回一个迭代器，生成索引和对应的值
func Enumerate(s []string) func(yield func(int, string) bool) {
    return func(yield func(int, string) bool) {
        for i, v := range s {
            // 传递索引 i 和 值 v
            if !yield(i, v) {
                return
            }
        }
    }
}

使用该迭代器：

func main() {
    items := []string{"apple", "banana", "cherry"}
    for i, name := range Enumerate(items) {
        fmt.Printf("Index: %d, Name: %s\n", i, name)
        // 当索引为 1 时测试 break 功能
        if i == 1 {
            break
        }
    }
}

4. 掌握迭代终止逻辑

理解 yield 函数的返回行为至关重要，它负责处理 break、return 和 goto 等控制流指令。

以下流程图展示了 for range 循环期间，迭代器、yield 函数与循环体之间的交互逻辑：

注意以下规则：

• 当循环体执行完毕自然进入下一次迭代时，yield 返回 true。
• 当循环体内遇到 break 或包含 range 的函数执行 return 时，yield 返回 false。
• 切勿在迭代器函数中忽略 yield 的返回值。如果忽略，即使外部使用了 break，你的迭代器仍会继续运行，导致资源浪费或逻辑错误。

5. 处理错误与资源清理

在自定义迭代器中，特别是涉及文件或网络连接时，必须确保在循环提前终止时能够正确清理资源。

使用 defer 语句是处理此问题的最佳实践。

1. 打开资源（如文件、数据库连接）。
2. 使用 defer 确保资源最终被关闭。
3. 检查 yield 的返回值，若为 false，执行 return 以触发 defer。

代码示例：

// Lines 按行读取文件内容
func Lines(path string) func(yield func(string) bool) {
    return func(yield func(string) bool) {
        // 打开文件
        f, err := os.Open(path)
        if err != nil {
            return // 如果打开失败，直接返回，yield 不会被调用，range 不会开始
        }
        // 确保文件关闭，无论循环是正常结束还是被 break 终止
        defer f.Close()

        scanner := bufio.NewScanner(f)
        for scanner.Scan() {
            line := scanner.Text()
            // 如果 yield 返回 false，通过 return 跳出循环，进而触发 defer f.Close()
            if !yield(line) {
                return
            }
        }
    }
}

6. 常见签名模式速查

下表总结了 Go 语言中自定义迭代器最常用的函数签名模式及其用途。

7. 泛型迭代器的实现

结合 Go 泛型，编写一个通用的迭代器可以适用于多种类型。

1. 定义泛型类型参数 T。
2. 使用 T 作为 yield 的参数类型。
3. 实现通用的遍历逻辑。

代码示例：

// All 是一个通用的切片迭代器
func All[T any](s []T) func(yield func(T) bool) {
    return func(yield func(T) bool) {
        for _, v := range s {
            if !yield(v) {
                return
            }
        }
    }
}

调用泛型迭代器时，类型通常可以被自动推导：

func main() {
    nums := []int{1, 2, 3}
    strs := []string{"a", "b", "c"}

    for n := range All(nums) {
        fmt.Println(n)
    }

    for s := range All(strs) {
        fmt.Println(s)
    }
}