golang

Развлекаемся с итераторами в Go

  • четверг, 24 октября 2024 г. в 00:00:06
https://habr.com/ru/companies/beget/articles/852940/

Релиз версии Go 1.23 добавил поддержку итераторов и пакет iter. Теперь можно перебирать константы, контейнеры (map, slice, array, string) и функции. Сначала создание итератора показалось мне неудобным, хотя в то же время его использование выглядело простым.

Моя проблема с подходом к итераторам в Go заключается в том, что их нельзя «связывать» так,как это можно делать в JavaScript:

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

Раздражение

Написание аналогичной конструкции на Go потребует цепочки из 5 вызовов функций:

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

Это пример, в пакете slices нет функций Map, Filter или ForEach.

Решение (вроде бы)

Поскольку я испытываю сильную неприязнь к написанию цепочек “функциональных” операций подобным образом (смотрю на тебя, Python, не набрасывайтесь на меня, хаскельщики), я хотел использовать новые итераторы и пакет iter, обернув их в структуру, которая позволяла бы писать чистую и аккуратную цепочку операций, как это реализовано в JavaScript.

Ниже приведены те же операции, но вместо использования пакетов iter и slices я использую свою абстракцию:

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

Логика

Давайте взглянем на реализацию, представляю вам структуру Iterator. Она оборачивает итератор (*Iterator).iter, что позволяет вызывать функции этой структуры, вместо передачи каждой функции итератора в качестве параметра следующей.

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

Давайте взглянем на первые функции, которые приходят на ум, когда мы говорим об итераторах: создание итератора из слайса и сбор его обратно в слайс:

func (i Iterator[V]) Collect() []V {
	collect := make([]V, 0)
	for e := range i.iter {
		collect = append(collect, e)
	}
	return collect
}

func From[V any](slice []V) *Iterator[V] {
	return &Iterator[V]{
		iter: func(yield func(V) bool) {
			for _, v := range slice {
				if !yield(v) {
					return
				}
			}
		},
	}
}

Первая функция максимально проста – создаем слайс, используем итератор, добавляем каждый элемент и возвращаем слайс. Вторая подчеркивает странный способ создания итераторов в Go. Давайте сначала посмотрим на сигнатуру: мы возвращаем указатель на структуру, чтобы вызывающий код мог вызывать все методы без необходимости использовать временную переменную для каждого вызова. В самой функции итератор создается путем возврата замыкания, которое выполняет цикл по параметру и возвращает результат, который останавливает итератор, когда функция yield возвращает false.

Each

Следующий метод, который я хочу реализовать – ForEach / Each. Он просто вызывает переданную функцию для каждого элемента итератора.

func (i *Iterator[V]) Each(f func(V)) {
	for i := range i.iter {
		f(i)
	}
}

Пример использования:

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

Reverse

Способ получить обратный итератор: сначала нужно собрать все элементы, а затем создать новый итератор из собранного слайса. К счастью, у нас есть функции, которые делают именно это:

func (i *Iterator[V]) Reverse() *Iterator[V] {
	collect := i.Collect()
	counter := len(collect) - 1
	for e := range i.iter {
		collect[counter] = e
		counter--
	}
	return From(collect)
}

Пример использования:

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

Map

Мутирование каждого элемента итератора также необходимо:

func (i *Iterator[V]) Map(f func(V) V) *Iterator[V] {
	cpy := i.iter
	i.iter = func(yield func(V) bool) {
		for v := range cpy {
			v = f(v)
			if !yield(v) {
				return
			}
		}
	}
	return i
}

Сначала мы копируем предыдущий итератор. Делая это, мы избегаем переполнения стека, ссылаясь на итератор i.iter в самом итераторе. Метод Map работает, переписывая i.iter новым итератором, который обрабатывает каждое поле копии итератора и заменяет значение итератора результатом передачи v в f, таким образом осуществляя отображение по итератору.

Filter

После Map, возможно, самым часто используемым методом функционального API является Filter. Давайте взглянем на нашу последнюю операцию:

func (i *Iterator[V]) Filter(f func(V) bool) *Iterator[V] {
	cpy := i.iter
	i.iter = func(yield func(V) bool) {
		for v := range cpy {
			if f(v) {
				if !yield(v) {
					return
				}
			}
		}
	}
	return i
}

Аналогично Map, мы копируем итератор и вызываем f с v в качестве параметра для каждого элемента. Если f возвращает true, мы сохраняем элемент в новом итераторе.

Примеры и мысли

slices и пакет iter отлично работают вместе с системой дженериков, введенной в Go 1.18.

Хотя этот вариант API прощен в использовании, я понимаю почему команда Go реализовала итераторы по другому. Ниже приведены тесты, которые служат примерами, и результаты их выполнения.

package iter1

import (
	"fmt"
	"testing"
	"unicode"
)

func TestIteratorNumbers(t *testing.T) {
	From([]int{1, 2, 3, 4}).
		Reverse().
		Map(func(i int) int { return i * i }).
		Filter(func(i int) bool { return i%2 == 0 }).
		Each(func(a int) { println(a) })
}

func TestIteratorRunes(t *testing.T) {
	r := From([]rune("Hello World!")).
		Reverse().
		// remove all spaces
		Filter(func(r rune) bool { return !unicode.IsSpace(r) }).
		// convert every rune to uppercase
		Map(func(r rune) rune { return unicode.ToUpper(r) }).
		Collect()
	fmt.Println(string(r))
}

func TestIteratorStructs(t *testing.T) {
	type User struct {
		Id   int
		Name string
		Hash int
	}

	u := []User{
		{0, "xnacly", 0},
		{1, "hans", 0},
		{2, "gedigedagedeio", 0},
	}

	From(u).
		// computing the hash for each user
		Map(func(u User) User {
			h := 0
			for i, r := range u.Name {
				h += int(r)*31 ^ (len(u.Name) - i - 1)
			}
			u.Hash = h
			return u
		}).
		Each(func(u User) { fmt.Printf("%#+v\n", u) })
}

Результаты запуска:

$ go test ./... -v
=== RUN   TestIteratorNumbers
16
4
--- PASS: TestIteratorNumbers (0.00s)
=== RUN   TestIteratorRunes
!DLROWOLLEH
--- PASS: TestIteratorRunes (0.00s)
=== RUN   TestIteratorStructs
&iter1.User{Id:0, Name:"xnacly", Hash:20314}
&iter1.User{Id:1, Name:"hans", Hash:13208}
&iter1.User{Id:2, Name:"gedigedagedeio", Hash:44336}
--- PASS: TestIteratorStructs (0.00s)
PASS
ok      iter1   0.263s

Вот и все, обертка в стиле JavaScript над iter и slices, готова.