// Цена цикла по функции: три группы, и внутри каждой строки делают ОДНУ И ТУ
// ЖЕ работу.
//
// ЧТО ЗДЕСЬ СРАВНИМО. Внутри группы — всё: каждая строка складывает одни и те
// же числа одного и того же среза и обязана дать одну и ту же сумму (это
// проверяется, см. TestSameWork). Отличается только способ обхода. Между
// группами сравнивать нельзя.
//
// ПОЧЕМУ СУММА, А НЕ ПУСТОЕ ТЕЛО. Пустое тело компилятор вправе выбросить
// вместе с циклом, и замер выродится в измерение пустоты. Сумма дешева, но
// наблюдаема: результат уходит в пакетную переменную, и выбросить его нельзя.
//
// ПОЧЕМУ ЦИКЛ ПО b.N, А НЕ b.Loop. Это здесь не стилистика, а условие
// осмысленности замера. Начиная с Go 1.24 компилятор НЕ ВКЛАДЫВАЕТ вызовы
// внутри тела b.Loop: в cmd/compile/internal/inline/interleaved/interleaved.go
// это записано прямым текстом — «No inlining nor devirtualization performed on
// b.Loop body». Для обычного кода это неважно, а для итератора важно как
// ничто другое: вся его цена и определяется тем, вложил компилятор функцию
// или нет. Замер на b.Loop измеряет только второй случай и о первом молчит.
// Обе стороны этой разницы измерены здесь отдельно — см. группу Noinline и
// одноимённые бенчмарки с суффиксом Loop.
//
// ЗАПУСК:
//
//	cd bench/goiter && ./cost.sh
//
// Снято на go1.24.7 linux/amd64.
package goiter

import (
	"iter"
	"slices"
	"testing"
)

const N = 10000

// Приёмник результата: не даёт компилятору выбросить тело цикла.
var sink int

func data() []int {
	s := make([]int, N)
	for i := range s {
		s[i] = i
	}
	return s
}

// Рукописный итератор — ровно то, что показано в документации пакета iter.
func values(s []int) iter.Seq[int] {
	return func(yield func(int) bool) {
		for _, v := range s {
			if !yield(v) {
				return
			}
		}
	}
}

// Обход обратным вызовом — способ, которым это писали до Go 1.23. Отличается
// от итератора ровно одним: тело не может ни выйти из внешней функции, ни
// прервать обход, потому что возвращать ему нечего.
func each(s []int, f func(int)) {
	for _, v := range s {
		f(v)
	}
}

// Итератор, пришедший параметром. Компилятор вправе вложить и sumSeq, и
// пришедший в него итератор — и вкладывает.
func sumSeq(seq iter.Seq[int]) int {
	sum := 0
	for v := range seq {
		sum += v
	}
	return sum
}

// То же самое, но вложить нельзя. Это не искусственный случай: ровно так
// выглядит любая функция, которая принимает iter.Seq и достаточно велика,
// чтобы не поместиться в бюджет вкладывания.
//
//go:noinline
func sumSeqNoinline(seq iter.Seq[int]) int {
	sum := 0
	for v := range seq {
		sum += v
	}
	return sum
}

// Итератор, спрятанный за пакетной переменной: какая именно функция там
// лежит, компилятор знать не может.
var storedSeq iter.Seq[int]

// ------------------------------------------- группа 1: компилятор видит всё

func BenchmarkPlain(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for _, v := range s {
			sum += v
		}
		sink = sum
	}
}

func BenchmarkSlicesValues(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for v := range slices.Values(s) {
			sum += v
		}
		sink = sum
	}
}

func BenchmarkHandwritten(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for v := range values(s) {
			sum += v
		}
		sink = sum
	}
}

func BenchmarkCallback(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		each(s, func(v int) { sum += v })
		sink = sum
	}
}

func BenchmarkSeqParam(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sink = sumSeq(slices.Values(s))
	}
}

// ------------------------------------ группа 2: компилятор не видит функцию

func BenchmarkNoinlinePlain(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for _, v := range s {
			sum += v
		}
		sink = sum
	}
}

func BenchmarkNoinlineSeqParam(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sink = sumSeqNoinline(slices.Values(s))
	}
}

func BenchmarkNoinlineStoredSeq(b *testing.B) {
	s := data()
	storedSeq = slices.Values(s)
	for i := 0; i < b.N; i++ {
		sum := 0
		for v := range storedSeq {
			sum += v
		}
		sink = sum
	}
}

// --------------------------------------------- группа 3: пара ключ-значение

func BenchmarkPairPlain(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for i, v := range s {
			sum += i * v
		}
		sink = sum
	}
}

func BenchmarkPairSlicesAll(b *testing.B) {
	s := data()
	for i := 0; i < b.N; i++ {
		sum := 0
		for i, v := range slices.All(s) {
			sum += i * v
		}
		sink = sum
	}
}

// ------------------------------------------------------- ловушка с b.Loop
//
// Те же две строки, что в группе 1, только на b.Loop. Разница между ними и
// группой 1 — не свойство итератора, а свойство замера: внутри b.Loop
// вкладывание выключено.

func BenchmarkLoopPlain(b *testing.B) {
	s := data()
	for b.Loop() {
		sum := 0
		for _, v := range s {
			sum += v
		}
		sink = sum
	}
}

func BenchmarkLoopSlicesValues(b *testing.B) {
	s := data()
	for b.Loop() {
		sum := 0
		for v := range slices.Values(s) {
			sum += v
		}
		sink = sum
	}
}

// --------------------------------------------------------------- проверка

// Все способы обхода обязаны давать одну и ту же сумму.
//
// Без этого замер сравнивал бы разную работу и молчал об этом: строка,
// случайно обходящая половину среза, выглядела бы просто быстрой.
func TestSameWork(t *testing.T) {
	s := data()
	want := 0
	for _, v := range s {
		want += v
	}

	cases := map[string]func() int{
		"slices.Values": func() int {
			got := 0
			for v := range slices.Values(s) {
				got += v
			}
			return got
		},
		"рукописный итератор": func() int {
			got := 0
			for v := range values(s) {
				got += v
			}
			return got
		},
		"обратный вызов": func() int {
			got := 0
			each(s, func(v int) { got += v })
			return got
		},
		"итератор параметром":           func() int { return sumSeq(slices.Values(s)) },
		"итератор параметром, noinline": func() int { return sumSeqNoinline(slices.Values(s)) },
	}
	for name, f := range cases {
		if got := f(); got != want {
			t.Errorf("%s: %d, ожидалось %d", name, got, want)
		}
	}

	storedSeq = slices.Values(s)
	got := 0
	for v := range storedSeq {
		got += v
	}
	if got != want {
		t.Errorf("итератор в пакетной переменной: %d, ожидалось %d", got, want)
	}

	wantPair := 0
	for i, v := range s {
		wantPair += i * v
	}
	gotPair := 0
	for i, v := range slices.All(s) {
		gotPair += i * v
	}
	if gotPair != wantPair {
		t.Errorf("slices.All: %d, ожидалось %d", gotPair, wantPair)
	}
}