stream.go

package stream

import (
	"bytes"
	"encoding/gob"

	"github.com/khaosles/giz/slice"
	"golang.org/x/exp/constraints"
)

/*
   @File: Stream.go
   @Author: khaosles
   @Time: 2023/8/13 13:38
   @Desc:
*/

// A StreamI should implements methods:
//type StreamI[T any] interface {
//
//	// part methods of Java Stream Specification.
//	Distinct() StreamI[T]
//	Filter(predicate func(item T) bool) StreamI[T]
//	FlatMap(mapper func(item T) StreamI[T]) StreamI[T]
//	Map(mapper func(item T) T) StreamI[T]
//	Peek(consumer func(item T)) StreamI[T]
//
//	Sorted(less func(a, b T) bool) StreamI[T]
//	Max(less func(a, b T) bool) (T, bool)
//	Min(less func(a, b T) bool) (T, bool)
//
//	Limit(maxSize int) StreamI[T]
//	Skip(n int) StreamI[T]
//
//	AllMatch(predicate func(item T) bool) bool
//	AnyMatch(predicate func(item T) bool) bool
//	NoneMatch(predicate func(item T) bool) bool
//	ForEach(consumer func(item T))
//	Reduce(init T, accumulator func(a, b T) T) T
//	Count() int
//
//	FindFirst() (T, bool)
//
//	ToSlice() []T
//
//	// part of methods custom extension
//	Reverse() StreamI[T]
//	Range(start, end int) StreamI[T]
//	Concat(Streams ...StreamI[T]) StreamI[T]
//}

type Stream[T any] struct {
	source []T
}

// Of creates a Stream whose elements are the specified values.
func Of[T any](elems ...T) Stream[T] {
	return FromSlice(elems)
}

// Generate Stream where each element is generated by the provided generater function
func Generate[T any](generator func() func() (item T, ok bool)) Stream[T] {
	source := make([]T, 0)

	var zeroValue T
	for next, item, ok := generator(), zeroValue, true; ok; {

		item, ok = next()
		if ok {
			source = append(source, item)
		}
	}

	return FromSlice(source)
}

// FromSlice creates Stream from slice.
func FromSlice[T any](source []T) Stream[T] {
	return Stream[T]{source: source}
}

// FromChannel creates Stream from channel.
func FromChannel[T any](source <-chan T) Stream[T] {
	s := make([]T, 0)

	for v := range source {
		s = append(s, v)
	}

	return FromSlice(s)
}

// FromRange creates a number Stream from start to end. both start and end are included. [start, end]
func FromRange[T constraints.Integer | constraints.Float](start, end, step T) Stream[T] {
	if end < start {
		panic("Stream.FromRange: param start should be before param end")
	} else if step <= 0 {
		panic("Stream.FromRange: param step should be positive")
	}

	l := int((end-start)/step) + 1
	source := make([]T, l)

	for i := 0; i < l; i++ {
		source[i] = start + (T(i) * step)
	}

	return FromSlice(source)
}

// Concat creates a lazily concatenated Stream whose elements are all the elements of the first Stream followed by all the elements of the second Stream.
func Concat[T any](a, b Stream[T]) Stream[T] {
	source := make([]T, 0)

	source = append(source, a.source...)
	source = append(source, b.source...)

	return FromSlice(source)
}

// Distinct returns a Stream that removes the duplicated items.
func (s Stream[T]) Distinct() Stream[T] {
	source := make([]T, 0)

	distinct := map[string]bool{}

	for _, v := range s.source {
		k := hashKey(v)
		if _, ok := distinct[k]; !ok {
			distinct[k] = true
			source = append(source, v)
		}
	}

	return FromSlice(source)
}

func hashKey(data any) string {
	buffer := bytes.NewBuffer(nil)
	encoder := gob.NewEncoder(buffer)
	err := encoder.Encode(data)
	if err != nil {
		panic("Stream.hashKey: get hashkey failed")
	}
	return buffer.String()
}

// Filter returns a Stream consisting of the elements of this Stream that match the given predicate.
func (s Stream[T]) Filter(predicate func(item T) bool) Stream[T] {
	source := make([]T, 0)

	for _, v := range s.source {
		if predicate(v) {
			source = append(source, v)
		}
	}

	return FromSlice(source)
}

// Map returns a Stream consisting of the elements of this Stream that apply the given function to elements of Stream.
func (s Stream[T]) Map(mapper func(item T) T) Stream[T] {
	source := make([]T, s.Count())

	for i, v := range s.source {
		source[i] = mapper(v)
	}

	return FromSlice(source)
}

// Peek returns a Stream consisting of the elements of this Stream, additionally performing the provided action on each element as elements are consumed from the resulting Stream.
func (s Stream[T]) Peek(consumer func(item T)) Stream[T] {
	for _, v := range s.source {
		consumer(v)
	}

	return s
}

// Skip returns a Stream consisting of the remaining elements of this Stream after discarding the first n elements of the Stream.
// If this Stream contains fewer than n elements then an empty Stream will be returned.
func (s Stream[T]) Skip(n int) Stream[T] {
	if n <= 0 {
		return s
	}

	source := make([]T, 0)
	l := len(s.source)

	if n > l {
		return FromSlice(source)
	}

	for i := n; i < l; i++ {
		source = append(source, s.source[i])
	}

	return FromSlice(source)
}

// Limit returns a Stream consisting of the elements of this Stream, truncated to be no longer than maxSize in length.
func (s Stream[T]) Limit(maxSize int) Stream[T] {
	if s.source == nil {
		return s
	}

	if maxSize < 0 {
		return FromSlice([]T{})
	}

	source := make([]T, 0, maxSize)

	for i := 0; i < len(s.source) && i < maxSize; i++ {
		source = append(source, s.source[i])
	}

	return FromSlice(source)
}

// AllMatch returns whether all elements of this Stream match the provided predicate.
func (s Stream[T]) AllMatch(predicate func(item T) bool) bool {
	for _, v := range s.source {
		if !predicate(v) {
			return false
		}
	}

	return true
}

// AnyMatch returns whether any elements of this Stream match the provided predicate.
func (s Stream[T]) AnyMatch(predicate func(item T) bool) bool {
	for _, v := range s.source {
		if predicate(v) {
			return true
		}
	}

	return false
}

// NoneMatch returns whether no elements of this Stream match the provided predicate.
func (s Stream[T]) NoneMatch(predicate func(item T) bool) bool {
	return !s.AnyMatch(predicate)
}

// ForEach performs an action for each element of this Stream.
func (s Stream[T]) ForEach(action func(item T)) {
	for _, v := range s.source {
		action(v)
	}
}

// Reduce performs a reduction on the elements of this Stream, using an associative accumulation function, and returns an Optional describing the reduced value, if any.
func (s Stream[T]) Reduce(initial T, accumulator func(a, b T) T) T {
	for _, v := range s.source {
		initial = accumulator(initial, v)
	}

	return initial
}

// Count returns the count of elements in the Stream.
func (s Stream[T]) Count() int {
	return len(s.source)
}

// FindFirst returns the first element of this Stream and true, or zero value and false if the Stream is empty.
func (s Stream[T]) FindFirst() (T, bool) {
	var result T

	if s.source == nil || len(s.source) == 0 {
		return result, false
	}

	return s.source[0], true
}

// FindLast returns the last element of this Stream and true, or zero value and false if the Stream is empty.
func (s Stream[T]) FindLast() (T, bool) {
	var result T

	if s.source == nil || len(s.source) == 0 {
		return result, false
	}

	return s.source[len(s.source)-1], true
}

// Reverse returns a Stream whose elements are reverse order of given Stream.
func (s Stream[T]) Reverse() Stream[T] {
	l := len(s.source)
	source := make([]T, l)

	for i := 0; i < l; i++ {
		source[i] = s.source[l-1-i]
	}
	return FromSlice(source)
}

// Range returns a Stream whose elements are in the range from start(included) to end(excluded) original Stream.
func (s Stream[T]) Range(start, end int) Stream[T] {
	if start < 0 {
		start = 0
	}
	if end < 0 {
		end = 0
	}
	if start >= end {
		return FromSlice([]T{})
	}

	source := make([]T, 0)

	if end > len(s.source) {
		end = len(s.source)
	}

	for i := start; i < end; i++ {
		source = append(source, s.source[i])
	}

	return FromSlice(source)
}

// Sorted returns a Stream consisting of the elements of this Stream, sorted according to the provided less function.
func (s Stream[T]) Sorted(less func(a, b T) bool) Stream[T] {
	source := []T{}
	source = append(source, s.source...)

	slice.SortBy(source, less)

	return FromSlice(source)
}

// Max returns the maximum element of this Stream according to the provided less function.
// less: a > b
func (s Stream[T]) Max(less func(a, b T) bool) (T, bool) {
	var max T

	if len(s.source) == 0 {
		return max, false
	}

	for i, v := range s.source {
		if less(v, max) || i == 0 {
			max = v
		}
	}
	return max, true
}

// Min returns the minimum element of this Stream according to the provided less function.
// less: a < b
func (s Stream[T]) Min(less func(a, b T) bool) (T, bool) {
	var min T

	if len(s.source) == 0 {
		return min, false
	}

	for i, v := range s.source {
		if less(v, min) || i == 0 {
			min = v
		}
	}

	return min, true
}

// ToSlice return the elements in the Stream.
func (s Stream[T]) ToSlice() []T {
	return s.source
}