Fun with Go's range over function types

A quick look at range-aware custom iterators.


Since Go 1.23, it’s possible to iterate over user-defined constructs with usual range loops, as long as such types conform to the expected interface.

Thanks to this new standard interface, it should be easier to integrate with third-party containers and to more concisely express some algorithms.

Range over integers

To start off, there’s been nice feature since 1.22 more or less on the same topic: A terser syntax for iterating over integers in the range [0, n). Similar to APL’s Iota primitive.

Example: Iterating from [0, 10)

for i := range 10 {
    fmt.Printf("i = %v\n", i)
}

Since the loop variable is optional, if we just want to execute an action multiple times, say 3 times, the code boils down to:

for range 3 {
    fmt.Println("some-action")
}

This is a nice little improvement in quality of life.

Range over user-defined containers

For illustration purposes, let’s build a Set that stores unique elements and ensure uniqueness by checking if an equal element has been already added before adding a new one:

import "slices"

type Set[V comparable] struct {
    values []V
}

func New[V comparable]() Set[V] {
    return Set[V]{}
}

func (s *Set[V]) Add(v V) {
    if s.Contains(v) {
        return
    }
    s.values = append(s.values, v)
}

func (s *Set[V]) Contains(v V) bool {
    return slices.Contains(s.values, v)
}

To make Set compatible with range, we define a method, here arbitrarily named Values(), that returns an iter.Seq from the new iter package:

import "iter"

func (s *Set[V]) Values() iter.Seq[V] {
    return func(yield func(V) bool) {
        for _, v := range s.values {
            if !yield(v) {
                return
            }
        }
    }
}

Values returns a function that itself accepts a function, typically – yet arbitrarily – named yield, that we call to push values into the loop.

Note that yield returns a boolean and that is false when the loop exits early, e.g. due to a break, and therefore we shouldn’t push any more values.

We use it like this:

s := set.New[int]()
s.Add(1)
s.Add(2)
s.Add(2)
s.Add(3)
for v := range s.Values() {
    fmt.Println(v)
}

Delegating to slices.Values

Within the same release, slices has received plenty of helper functions for iteration, including Values. Equipped with it, our Values reduces to a one-line delegation:

func (s *Set[V]) Values() iter.Seq[V] {
    return slices.Values(s.values)
}

Range over infinite sequences

We can also range over infinite sequences, for example, Fibonacci’s:

package main

func fibs() iter.Seq[int] {
    return func(yield func(int) bool) {
        a, b := 0, 1
        for {
            if !yield(a) {
                return
            }
            a, b = b, a+b
        }
    }
}

func main() {
    for f := range fibs() {
        if f > 100 {
            break
        }
        fmt.Println(f)
    }
}

Since the iterator returned by fibs produces an infinite sequence, we had to explicitly break out of the loop to prevent it from running “forever”.

Range adapters

Lastly, we can write functions that adapts iterators to our needs, for instance, takeWhile to keep iterating while a given predicate applied to the current value yields true. As an example, we use to iterate over an infinite sequence of strings (a aa aaa …) and take only the firsts whose lengths are < 5.

package main

func appendForever(s string) iter.Seq[string] {
    init := s
    return func(yield func(string) bool) {
        for {
            if !yield(s) {
                return
            }
            s = init + s
        }
    }
}

func takeWhile[V any](p func(V) bool, seq iter.Seq[V]) iter.Seq[V] {
    return func(yield func(V) bool) {
        for v := range seq {
            if !p(v) || !yield(v) {
                return
            }
        }
    }
}

func main() {
    for s := range takeWhile(func(s string) bool { return len(s) < 5 }, appendForever("a")) {
        fmt.Println(s)
    }
}

Conclusion

There’s a lot to discuss about for this lovely feature, for example the contrast between push iterators (all examples shown here) and pull iterators, where the control-flow is more or less reversed. We could also write different methods for different iterations, e.g. forwards, backwards, in-order, pre-order, etc.

I encourage you to go through Range Over Function Types for more details.

Tags: go
Share: X (Twitter) Facebook LinkedIn