Golang Rangefunc

Golang 1.22 中加入了实验性功能 Rangefunc，使用环境变量 GOEXPERIMENT=rangefunc 来开启。

Updates

自 Golang 1.23 起，Rangefunc 的功能已经加入了标准库中，无需使用环境变量来开启

开启之后，就可以引用 iter 包，包含两个 type 和两个 helper functions.

type Seq[V any] func(yield func(V) bool)
type Seq2[K, V any] func(yield func(K, V) bool)
 
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())

具体的使用方法可以参考下面这个简单的例子。

package main
 
import "fmt"
import "iter"
 
func backwards[E any](s []E) iter.Seq[E] {
    return func(yield func(E) bool) {
        for i := len(s) - 1; i >= 0; i-- {
            if !yield(s[i]) {
                return
            }
        }
    }
 
}
 
func main() {
    for x := range backwards([]string{"hello", "world"}) {
        fmt.Println(x)
    }
}

方法 iter.Pull 可以从一个 iter 对象中拉取数据，比如下面的 zip 函数。

func zip[T1, T2 any](seq1 iter.Seq[T1], seq2 iter.Seq[T2]) iter.Seq2[T1, T2] {
	return func(yield func(T1, T2) bool) {
		next1, stop1 := iter.Pull(seq1)
		next2, stop2 := iter.Pull(seq2)
		defer stop1()
		defer stop2()
 
		for {
			v1, flg1 := next1()
			v2, flg2 := next2()
			if !flg1 || !flg2 || !yield(v1, v2) {
				return
			}
		}
	}
}

在实现上，Golang 实际是通过编译时 rewrite 的方法来实现的对 iter 的支持，在第一个例子中实际的代码会被 rewrite 为下面的形式。

backward(s)(func(x string) bool {
    fmt.Println(x)
    return true
})

这里的 return false 就相当于 break，当然有其他复杂情况需要处理，完整的 rewrite 规则可以参考源码中的注释。

另外虽然是 rewrite 为 function 来实现，但是在 range loop 中调用 defer 也是和普通的 loop 相同，即在包含该 loop 的 function 结束时执行。

例如下面这个例子：

package main
 
import (
	"fmt"
	"iter"
)
 
func DeferAfterYield() iter.Seq[int] {
	defer fmt.Println("defer before range function")
	return func(yield func(int) bool) {
		defer fmt.Println("defer in range function")
 
		for i := 0; i < 3; i++ {
			yield(i)
		}
	}
}
 
func main() {
	for x := range DeferAfterYield() {
		defer func() {
			fmt.Println("defer in main loop: ", x)
		}()
 
		fmt.Println(x)
	}
	fmt.Println("other operation in main")
	defer func() {
		fmt.Println("defer in main")
	}()
}

其运行的结果为

defer before range function
0
1
2
defer in range function
other operation in main
defer in main
defer in main loop:  2
defer in main loop:  1
defer in main loop:  0

通过这个例子可以发现：

1. defer in range function 会在 loop 结束时调用
2. defer in main loop 是在 main 函数结束时调用，同一般 loop 中使用 defer
3. 变量 x 在每次循环中会生成新的，闭包函数可正常的捕获，不再需要想以前需要使用 var x1 := x 来捕获，这个不是 rangefunc 中提供的功能，而是新版本的 golang 修改的，这里只是提醒一下
4. 再提醒下 defer 函数的执行顺序，和栈逻辑一致，后进先出

说了这么多，那么 rangefunc 可以用在哪些地方呢？以下是几个例子：

和 python 类似的 range 函数

func Range(start, end, step int) iter.Seq[int] {
	if step == 0 {
		step = 1
	}
 
	return func(yield func(int) bool) {
		if step > 0 {
			for i := start; i < end; i += step {
				if !yield(i) {
					return
				}
			}
 
		} else {
			for i := start; i > end; i += step {
				if !yield(i) {
					return
				}
			}
		}
	}
}

用于数据库查询

func BqQuery[E any](ctx context.Context, query *bigquery.Query) iter.Seq2[*E, error] {
	return func(yield func(*E, error) bool) {
		iter, err := query.Read(ctx)
		if err != nil {
			yield(nil, err)
			return
		}
 
		for {
			var row E
			err := iter.Next(&row)
			if err != nil {
				if err != iterator.Done {
					// call error handler
					if !yield(nil, err) {
						return
					}
				}
				return
			}
 
			if !yield(&row, nil) {
				return
			}
		}
	}
}

按行读取文件

func Readlines(file string) iter.Seq2[[]byte, error] {
	return func(yield func([]byte, error) bool) {
		f, err := os.Open(file)
		if err != nil {
			yield(nil, err)
			return
		}
		defer f.Close()
 
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			if !yield(scanner.Bytes(), nil) {
				return
			}
		}
 
		if err := scanner.Err(); err != nil {
			yield(nil, err)
		}
	}
}

参考连接：

https://blog.perfects.engineering/go_range_over_funcs