sequence.go

package goiter

import (
    "iter"
    "math"
)

type GeneratorFunc[T any] func() (T, bool)

type GeneratorFunc2[T1, T2 any] func() (T1, T2, bool)

type TInt interface {
    ~int | ~int8 | ~int16 | ~int32 | ~int64 |
    ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64
}

// Range returns an iterator that yields a sequence of integers forward or backward from start to end, incrementing/decrementing by 1.
// to be specific, the second parameter "end" is inclusive.
// for example:
//
//  goiter.Range(1, 5) // will yield 1, 2, 3, 4, 5
//  goiter.Range(3, -3) // will yield 3, 2, 1, 0, -1, -2, -3
func Range[T TInt](start, end T) Iterator[T] {
    return RangeStep(start, end, 1)
}

// RangeStep extends the ability to Range function, allowing iteration from any integer and stepping forward or backward in any step.
// It is similar to Python's range function, but with some differences:
//  1. the second parameter "end" is inclusive, rather than exclusive in Python's range function.
//  2. stepSize does not accept negative numbers. Whether iterating in increasing or decreasing order, stepSize must be positive,
//     so you don't need to consider adjusting the sign of step according to the direction of iteration, you can consider it as the absolute value of the step parameter of Python range function.
//  3. Providing a value less than or equal to 0 for stepSize will not return an error, it simply doesn't yield any values.
func RangeStep[T TInt, S TInt](start, end T, stepSize S) Iterator[T] {
    if stepSize <= 0 {
        // 0 will lead to infinite loops
        return Empty[T]()
    }

    step := uint64(stepSize)
    inc := true
    if start > end {
        inc = false
    }

    if willOverflow(start, step, inc) {
        return func(yield func(T) bool) {
            yield(start)
        }
    }

    return func(yield func(T) bool) {
        curr := start
        for {
            if !yield(curr) {
                return
            }

            if inc {
                next := curr + T(step)
                if next > end || next < start || next <= curr {
                    // both next < start and next <= curr means overflow
                    return
                }
                curr = next
            } else {
                next := curr - T(step)
                if next < end || next > start || next >= curr {
                    // both next > start and next >= curr means overflow
                    return
                }
                curr = next
            }
        }
    }
}

// Counter returns an iterator that yields a sequence of integers incrementing by 1.
func Counter(startFrom int) Iterator[int] {
    var next = startFrom
    return Sequence(func() (int, bool) {
        v := next
        next++
        return v, true
    })
}

// Sequence takes a generator function and returns an iterator that yields the values generated by the generator.
// This is a general sequence generator function
// For example, you can use it to generate the Fibonacci sequence like this:
//
//	genFib := func(n int) goiter.GeneratorFunc[int] {
//		a, b := 0, 1
//		return func() (int, bool) {
//			if n <= 0 {
//				return 0, false
//			}
//			n--
//			a, b = b, a+b
//			return a, true
//		}
//	}
//	// this will print first 5 Fibonacci numbers: 1 1 2 3 5
//	for v := range goiter.Sequence(genFib(5)) {
//		fmt.Printf("%d ", v)
//	}
func Sequence[T any](generator func() (T, bool)) Iterator[T] {
    return func(yield func(T) bool) {
        for {
            v, hasMore := generator()
            if !hasMore {
                return
            }
            if !yield(v) {
                return
            }
        }
    }
}

// Sequence2 is the iter.Seq2 version of Sequence function.
func Sequence2[T1, T2 any](generator func() (T1, T2, bool)) Iterator2[T1, T2] {
    return func(yield func(T1, T2) bool) {
        for {
            v1, v2, hasMore := generator()
            if !hasMore {
                return
            }
            if !yield(v1, v2) {
                return
            }
        }
    }
}

// Reverse returns an iterator that yields the values of the input iterator in reverse order.
// So if the input iterator yields "a" "b" "c", then goiter.Reverse(iterator) will yield "c" "b" "a".
//
// be careful, if this function is used on iterators that has massive amount of data, it might consume a lot of memory.
func Reverse[TIter SeqX[T], T any](iterator TIter) Iterator[T] {
    return func(yield func(T) bool) {
        var buffer []T
        next, stop := iter.Pull(iter.Seq[T](iterator))
        defer stop()
        for {
            v, ok := next()
            if !ok {
                break
            }
            buffer = append(buffer, v)
        }
        for i := len(buffer) - 1; i >= 0; i-- {
            if !yield(buffer[i]) {
                return
            }
        }
    }
}

// Reverse2 is the iter.Seq2 version of Reverse function.
// be careful, if this function is used on iterators that has massive amount of data, it might consume a lot of memory.
func Reverse2[TIter Seq2X[T1, T2], T1, T2 any](iterator TIter) Iterator2[T1, T2] {
    return func(yield func(T1, T2) bool) {
        var buffer []*Combined[T1, T2]
        next, stop := iter.Pull2(iter.Seq2[T1, T2](iterator))
        defer stop()
        for {
            v1, v2, ok := next()
            if !ok {
                break
            }
            buffer = append(buffer, &Combined[T1, T2]{V1: v1, V2: v2})
        }
        for i := len(buffer) - 1; i >= 0; i-- {
            if !yield(buffer[i].V1, buffer[i].V2) {
                return
            }
        }
    }
}

func willOverflow[T TInt](v T, step uint64, inc bool) bool {
    tMax := int64(tMax(v))
    tMin := int64(tMin(v))

    if tMax != math.MaxInt64 && step >= uint64(tMax-tMin+1) {
        return true
    }
    if inc && v+T(step) < v {
        return true
    }
    if !inc && v-T(step) > v {
        return true
    }

    return false
}

func tMin[T TInt](v T) T {
    ones := ^T(0)
    if ones < 0 {
        return ^(ones ^ (1 << (countBits(ones) - 1)))
    }
    return 0
}

func tMax[T TInt](v T) T {
    ones := ^T(0)
    if ones < 0 {
        return ones ^ (1 << (countBits(ones) - 1))
    }
    return ones
}

func countBits[T TInt](v T) int {
    v = 1
    for _, bits := range [4]int{8, 16, 32} {
        if v<<bits == 0 {
            return bits
        }
    }
    return 64
}