range_iters.go

package main

import (
	"fmt"
	"io"
	"iter"
	"math/big"
	"math/rand"
)

func iter0(yield func() bool) {
	for range 3 {
		if !yield() {
			return
		}
	}
}

func testFuncRange0() {
	for range iter0 {
		fmt.Println("iter0")
	}
}

func iter1(yield func(i int) bool) {
	for i := range 3 {
		if !yield(i) {
			return
		}
	}
}

func testFuncRange1() {
	for i := range iter1 {
		fmt.Println("iter1", i)
	}
}

func plusOneIter(yield func(i int) bool) {
	for i := range 3 {
		if !yield(i+1) {
			return
		}
	}
}

func testFuncRangePlusOne() {
	for i := range plusOneIter {
		fmt.Println("iter1", i)
	}
}

func brokenIter(yield func(i int) bool) {
	for i := range 3 {
		yield(i+1)
	}
}

func testBrokenIter() {
	for i := range brokenIter {
		fmt.Println("brokenIter", i)
		if i > 1 {
			break
		}
	}
}

func iter2(yield func(i int, e int) bool) {
	for i := range 3 {
		if !yield(i, i+1) {
			return
		}
	}
}

func testFuncRange2() {
	for i, e := range iter2 {
		fmt.Println("iter2", i, e)
	}
}

type Slice []int

func (s Slice) All() func(yield func(i int) bool) {
	return func(yield func(i int) bool) {
		for i := range s {
			if !yield(s[i]) {
				return
			}
		}
	}
}

func (s Slice) Primes() func (yield func(i int) bool) {
	return func (yield func(i int) bool) {
		for i := range s {
			if big.NewInt(int64(s[i])).ProbablyPrime(0) {
				if !yield(s[i]) {
					return
				}
			}
		}
	}
}

func (s Slice) Primes2() func (yield func(i int, e int) bool) {
	return func (yield func(i int, e int) bool) {
		for i := range s {
			if big.NewInt(int64(s[i])).ProbablyPrime(0) {
				if !yield(i, s[i]) {
					return
				}
			}
		}
	}
}

func (s Slice) ErrorIter() func(yield func(i int, e error) bool) {
	return func(yield func(i int, e error) bool) {
		for _, i := range s {
			if !yield(i, nil) {
				return
			}
		}
	}
}

func (s Slice) PrimeIndexes() func (yield func(i int) bool) {
	return func (yield func(i int) bool) {
		for i := range s {
			if big.NewInt(int64(s[i])).ProbablyPrime(0) {
				if !yield(i) {
					return
				}
			}
		}
	}
}

func (s Slice) FilteredIter(predicate func(i int) bool) func (yield func(i int) bool) {
	return func (yield func(i int) bool) {
		for i := range s {
			if predicate(s[i]) {
				if !yield(s[i]) {
					return
				}
			}
		}
	}
}

func iterEvens(slice Slice) {
	for i := range slice.FilteredIter(func(i int) bool {
		return i%2 == 0
	}) {
		fmt.Println("even number:", i)
		if i > 10 {
			break
		}
	}
}

func iterAll(slice Slice) {
	for i := range slice.All() {
		fmt.Println("all iter:", i)
		if i > 10 {
			break
		}
	}
}

func iterPrimes(slice Slice) {
	for i := range slice.Primes() {
		fmt.Println("prime number:", i)
		if i > 10 {
			break
		}
	}
}

func iterPrimes2(slice Slice) {
	for i, e := range slice.Primes2() {
		fmt.Printf("prime number: %d, %d\n", i, e)
		if i > 10 {
			break
		}
	}
}

func iterWithErr(slice Slice) error {
	for i, err := range slice.ErrorIter() {
		if err != nil {
			return err
		}
		fmt.Printf("error iter got value: %d\n", i)
		if i > 10 {
			break
		}
	}
	return nil
}

func iterPrimesIndex(slice Slice) {
	for i := range slice.PrimeIndexes() {
		fmt.Printf("prime number at index %d, %d\n", i, slice[i])
		if i > 10 {
			break
		}
	}
}

func iterTraditional(slice Slice) {
	iter := slice.Iter()
	for {
		next, err := iter.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			fmt.Printf("error: %s\n", err.Error())
		}
		
		fmt.Println("iter got value: ", next)
	}
}

func iterTraditionalWithRange(slice Slice) {
	for i, err := range slice.RangeCompatibleIter() {
		if err != nil {
			fmt.Printf("error: %s\n", err.Error())
		}

		fmt.Println("iter got value: ", i)
	}
}

func iterTraditionalWithRangeRoundTrip(slice Slice) {
	next, stop := iter.Pull2(slice.RangeCompatibleIter())
	defer stop()
	i := 0
	for {
		result, err, valid := next()
		if !valid {
			break
		}
		
		if i > 10 {
			break
		}
		i++

		if err != nil {
			fmt.Printf("error: %s\n", err.Error())
		} else {
			fmt.Println("iter got value: ", result)
		}
	}
}

type iterEof struct {
	slice Slice
	i int
}

func (iter *iterEof) Next() (int, error) {
	defer func() {
		iter.i++
	}()
	
	if iter.i >= len(iter.slice) {
		return 0, io.EOF
	} else if rand.Float32() > .9 {
		return 0, fmt.Errorf("failed to fetch next element")
	}
	
	return iter.slice[iter.i], nil
}

func(s Slice) Iter() *iterEof {
	return &iterEof{slice: s}
}

func (s Slice) RangeCompatibleIter() func (yield func(int, error) bool) {
	iter := s.Iter()
	return func (yield func(i int, e error) bool) {
		for {
			next, err := iter.Next()
			if err == io.EOF {
				return
			}

			if !yield(next, err) {
				return 
			}
		}
	}
}

func main() {
	testFuncRange0()
	testFuncRange1()
	testFuncRangePlusOne()
	// This panics, uncomment to see
	// testBrokenIter()
	testFuncRange2()
	
	slice := make(Slice, 20)
	for i := 0; i < 20; i++ {
		slice[i] = i
	}
	rand.Shuffle(len(slice), func(i, j int) {
		slice[i], slice[j] = slice[j], slice[i]
	})

	iterAll(slice)
	iterPrimes(slice)
	iterWithErr(slice)
	iterPrimes2(slice)
	iterPrimesIndex(slice)
	iterEvens(slice)

	iterTraditional(slice)
	iterTraditionalWithRange(slice)
	iterTraditionalWithRangeRoundTrip(slice)
}