Fun with Go Iterators

· 1207 words · 6 minute read

Go version 1.23 added iterator support 1 and the iter package 2. We can now loop over constants, containers (maps, slices, arrays, strings) and functions. At first I found the iterator creation clunky, while consuming the iterator seems straightforward.

My issue with the go way of iterators is, you can’t chain them like you would in JavaScript:

1[1,2,3,4]
2    .reverse()
3    .map(e => e*e)
4    .filter(e => e % 2 == 0)
5    .forEach(e => console.log(e)) 

The Annoyance ##

Writing this in go would require us to chain 5 function calls:

 1slices.ForEach(
 2    slices.Filter(
 3        slices.Map(
 4            slices.Reverse(slices.All([]int{1,2,3,4})), 
 5            func(i int) int { return i * i},
 6        ),
 7        func(i int) bool { return i % 2 == 0 }
 8    ),
 9    func(i int) { fmt.Println(i) }
10)

This is an example, there are are no Map, Filter or ForEach functions in the slices package 3.

The Solution (sort of) ##

Because i have a big distaste for writing chained “functional” operations like this, looking at you python (don’t come at me haskell bros) - I wanted to use the new iterators and the iter package and wrap it with a structure to allow the nice and clean chaining JavaScript provides. Below are the same operations, but instead of using the iter and slices package, I use my abstraction:

1func TestIterator(t *testing.T) {
2	From([]int{1, 2, 3, 4}).
3		Reverse().
4		Map(func(i int) int { return i * i }).
5		Filter(func(i int) bool { return i%2 == 0 }).
6		Each(func(a int) { println(a) })
7    // 16
8    // 4
9}

The Logic ##

Lets take a look a the implementation and let me introduce the Iteratorstruct. It wraps the iterator (*Iterator).iter and thus allows me to callfunctions on this structure, instead of having to use each iterator functionas a param to the next one.

1type Iterator[V any] struct {
2	iter iter.Seq[V]
3}

Lets take a look at the first functions coming to mind when talking about iterators, creating one from a slice and collection one into a slice:

 1func (i Iterator[V]) Collect() []V {
 2	collect := make([]V, 0)
 3	for e := range i.iter {
 4		collect = append(collect, e)
 5	}
 6	return collect
 7}
 8
 9func From[V any](slice []V) *Iterator[V] {
10	return &Iterator[V]{
11		iter: func(yield func(V) bool) {
12			for _, v := range slice {
13				if !yield(v) {
14					return
15				}
16			}
17		},
18	}
19}

The first one is as straight forward as possible - create a slice, consume the iterator, append each element, return the slice. The second highlights the weird way iterators are created in go. Lets first take a look at the signature, we are returning a pointer to the struct, so the callee can invoke all methods without having to use a temporary variable for each. In the function itself, the iterator is created by returning a closure, that loops over the param and returning, which stops the iterator, when the yield function returns false.

Each ###

The ForEach / Each method is the next function I want, It simply executes the passed in function for every element of the iterator.

1func (i *Iterator[V]) Each(f func(V)) {
2	for i := range i.iter {
3		f(i)
4	}
5}

It can be used like this:

1From([]int{1, 2, 3, 4}).Each(func(a int) { println(a) })
2// 1
3// 2
4// 3
5// 4

Reverse ###

A way to reverse the iterator, we first need to collect all elements and after that construct a new iterator from the collected slice, luckily we have functions for exactly this:

1func (i *Iterator[V]) Reverse() *Iterator[V] {
2	collect := i.Collect()
3	counter := len(collect) - 1
4	for e := range i.iter {
5		collect[counter] = e
6		counter--
7	}
8	return From(collect)
9}

Again, useful to reverse a slice:

1From([]int{1, 2, 3, 4}).Reverse().Each(func(a int) { println(a) })
2// 4
3// 3
4// 2
5// 1

Map ###

Mutating every element of the iterator is a necessity, too:

 1func (i *Iterator[V]) Map(f func(V) V) *Iterator[V] {
 2	cpy := i.iter
 3	i.iter = func(yield func(V) bool) {
 4		for v := range cpy {
 5			v = f(v)
 6			if !yield(v) {
 7				return
 8			}
 9		}
10	}
11	return i
12}

At first we copy the previous iterator, by doing this, we skip causing a stack overflow by referencing the i.iter iterator in the iterator itself. Map is works by overwriting the i.iter with a new iterator thats consuming every field of the cpy iterator and overwriting the iterator value with the result of passing v to f, thus mapping over the iterator.

Filter ###

After mapping, possibly the most used streaming / functional api method is Filter. So lets take a look at our final operation:

 1func (i *Iterator[V]) Filter(f func(V) bool) *Iterator[V] {
 2	cpy := i.iter
 3	i.iter = func(yield func(V) bool) {
 4		for v := range cpy {
 5			if f(v) {
 6				if !yield(v) {
 7					return
 8				}
 9			}
10		}
11	}
12	return i
13}

Similar to Map, we consume the copied iterator and invoke f with v as the param for each one, if f returns true, we keep it in the new iterator.

Examples and Thoughts ##

The slices and the iter package both play very good together with the generic system introduced in go 1.18 4.

While this API is easier to use, I understand the reasoning of the go team for not implementing iterators like this. Below are some tests that serve as examples and the result of running them.

 1package iter1
 2
 3import (
 4	"fmt"
 5	"testing"
 6	"unicode"
 7)
 8
 9func TestIteratorNumbers(t *testing.T) {
10	From([]int{1, 2, 3, 4}).
11		Reverse().
12		Map(func(i int) int { return i * i }).
13		Filter(func(i int) bool { return i%2 == 0 }).
14		Each(func(a int) { println(a) })
15}
16
17func TestIteratorRunes(t *testing.T) {
18	r := From([]rune("Hello World!")).
19		Reverse().
20		// remove all spaces
21		Filter(func(r rune) bool { return !unicode.IsSpace(r) }).
22		// convert every rune to uppercase
23		Map(func(r rune) rune { return unicode.ToUpper(r) }).
24		Collect()
25	fmt.Println(string(r))
26}
27
28func TestIteratorStructs(t *testing.T) {
29	type User struct {
30		Id   int
31		Name string
32		Hash int
33	}
34
35	u := []User{
36		{0, "xnacly", 0},
37		{1, "hans", 0},
38		{2, "gedigedagedeio", 0},
39	}
40
41	From(u).
42		// computing the hash for each user
43		Map(func(u User) User {
44			h := 0
45			for i, r := range u.Name {
46				h += int(r)*31 ^ (len(u.Name) - i - 1)
47			}
48			u.Hash = h
49			return u
50		}).
51		Each(func(u User) { fmt.Printf("%#+v\n", u) })
52}

Running these, results in:

 1$ go test ./... -v
 2=== RUN   TestIteratorNumbers
 316
 44
 5--- PASS: TestIteratorNumbers (0.00s)
 6=== RUN   TestIteratorRunes
 7!DLROWOLLEH
 8--- PASS: TestIteratorRunes (0.00s)
 9=== RUN   TestIteratorStructs
10&iter1.User{Id:0, Name:"xnacly", Hash:20314}
11&iter1.User{Id:1, Name:"hans", Hash:13208}
12&iter1.User{Id:2, Name:"gedigedagedeio", Hash:44336}
13--- PASS: TestIteratorStructs (0.00s)
14PASS
15ok      iter1   0.263s

So there u have it, a wrapper around iter and slices to mirror the way JavaScript provides streaming, only in go.