iterator.go

v0.6.0

package collect

import (
	"reflect"
	"slices"
)

// ListGenerator produces values one at a time. It returns (value, true) for
// each element and (zero, false) when the sequence is exhausted.
type ListGenerator[E any] func() (E, bool)

// ListIterator is a lazy, pull-based iterator over a sequence of elements.
// All intermediate operations (Map, Filter, Sort, Skip, Limit) return a new
// iterator; terminal operations (Count, Apply, Collect) consume it.
type ListIterator[E any] interface {
	// Map returns an iterator that applies f to each element.
	Map(f func(E) E) ListIterator[E]
	// Filter returns an iterator that yields only elements satisfying pred.
	Filter(pred func(E) bool) ListIterator[E]
	// Sort collects all elements, sorts them with less, and returns a new iterator.
	Sort(less func(E, E) bool) ListIterator[E]
	// Skip returns an iterator that discards the first count elements.
	// Panics if count is negative.
	Skip(count int) ListIterator[E]
	// Limit returns an iterator that yields at most count elements.
	// Panics if count is negative.
	Limit(count int) ListIterator[E]
	// Count exhausts the iterator and returns the number of elements.
	Count() int
	// Apply exhausts the iterator calling f for each element and returns the count.
	Apply(f func(E)) int
	// Collect collects all remaining elements into a slice.
	Collect() []E
	// MapIndex returns an iterator that applies f with the 0-based index to each element.
	MapIndex(f func(int, E) E) ListIterator[E]
	// FilterIndex returns an iterator that yields only elements where pred(index, elem) is true.
	FilterIndex(pred func(int, E) bool) ListIterator[E]
	// Unique returns an iterator that deduplicates elements using the string key returned by fn.
	Unique(fn func(E) string) ListIterator[E]
	// ApplyIndex exhausts the iterator calling f with the 0-based index for each element.
	ApplyIndex(f func(int, E))
	// Materialize materializes the iterator and returns a new iterator over the result.
	Materialize() ListIterator[E]
}

type listIterator[E any] struct {
	generator   ListGenerator[E]
	maxElements int
}

// NewListIterator creates a [ListIterator] backed by gen. maxElements is a
// hint for the expected upper bound of elements (used for pre-allocation);
// pass -1 if unknown.
func NewListIterator[E any](gen ListGenerator[E], maxElements int) ListIterator[E] {
	return &listIterator[E]{generator: gen, maxElements: maxElements}
}

// FromSlice creates a [ListIterator] that yields the elements of inp
// in order.
func FromSlice[E any](inp []E) ListIterator[E] {
	pos := 0
	gen := func() (E, bool) {
		if pos >= len(inp) {
			var zero E
			return zero, false
		}
		v := inp[pos]
		pos++
		return v, true
	}
	return &listIterator[E]{generator: gen, maxElements: len(inp)}
}

// FromSliceWith creates a [ListIterator] that maps each element of inp through cb.
func FromSliceWith[E any, I any](inp []E, cb func(E) I) ListIterator[I] {
	pos := 0
	gen := func() (I, bool) {
		if pos >= len(inp) {
			var zero I
			return zero, false
		}
		v := cb(inp[pos])
		pos++
		return v, true
	}
	return &listIterator[I]{generator: gen, maxElements: len(inp)}
}

// FromMapKeys creates a [ListIterator] over the keys of a map.
func FromMapKeys[K comparable, V any](inp map[K]V) ListIterator[K] {
	rv := reflect.ValueOf(inp)
	keys := rv.MapKeys()
	pos := 0
	gen := func() (K, bool) {
		if pos >= len(keys) {
			var zero K
			return zero, false
		}
		v := keys[pos].Interface().(K)
		pos++
		return v, true
	}
	return &listIterator[K]{generator: gen, maxElements: len(keys)}
}

// FromMapValues creates a [ListIterator] over the values of a map.
func FromMapValues[K comparable, V any](inp map[K]V) ListIterator[V] {
	rv := reflect.ValueOf(inp)
	keys := rv.MapKeys()
	pos := 0
	gen := func() (V, bool) {
		if pos >= len(keys) {
			var zero V
			return zero, false
		}
		v := rv.MapIndex(keys[pos]).Interface().(V)
		pos++
		return v, true
	}
	return &listIterator[V]{generator: gen, maxElements: len(keys)}
}

func (it *listIterator[E]) Map(f func(E) E) ListIterator[E] {
	gen := it.generator
	return &listIterator[E]{
		generator: func() (E, bool) {
			v, ok := gen()
			if !ok {
				return v, false
			}
			return f(v), true
		},
		maxElements: it.maxElements,
	}
}

func (it *listIterator[E]) Filter(pred func(E) bool) ListIterator[E] {
	gen := it.generator
	return &listIterator[E]{
		generator: func() (E, bool) {
			for {
				v, ok := gen()
				if !ok {
					return v, false
				}
				if pred(v) {
					return v, true
				}
			}
		},
		maxElements: it.maxElements,
	}
}

func (it *listIterator[E]) Sort(less func(E, E) bool) ListIterator[E] {
	all := it.Collect()
	slices.SortFunc(all, func(a, b E) int {
		switch {
		case less(a, b):
			return -1
		case less(b, a):
			return 1
		default:
			return 0
		}
	})
	return FromSlice(all)
}

func (it *listIterator[E]) Skip(count int) ListIterator[E] {
	if count < 0 {
		panic("collect: Skip count must be non-negative")
	}
	gen := it.generator
	skipped := false
	return &listIterator[E]{
		generator: func() (E, bool) {
			if !skipped {
				skipped = true
				for range count {
					if _, ok := gen(); !ok {
						var zero E
						return zero, false
					}
				}
			}
			return gen()
		},
		maxElements: max(it.maxElements-count, 0),
	}
}

func (it *listIterator[E]) Limit(count int) ListIterator[E] {
	if count < 0 {
		panic("collect: Limit count must be non-negative")
	}
	gen := it.generator
	remaining := count
	return &listIterator[E]{
		generator: func() (E, bool) {
			if remaining <= 0 {
				var zero E
				return zero, false
			}
			remaining--
			return gen()
		},
		maxElements: min(it.maxElements, count),
	}
}

func (it *listIterator[E]) Count() int {
	n := 0
	for {
		if _, ok := it.generator(); !ok {
			return n
		}
		n++
	}
}

func (it *listIterator[E]) Apply(f func(E)) int {
	n := 0
	for {
		v, ok := it.generator()
		if !ok {
			return n
		}
		f(v)
		n++
	}
}

func (it *listIterator[E]) Collect() []E {
	sz := max(it.maxElements, 0)
	rv := make([]E, 0, sz)
	for {
		v, ok := it.generator()
		if !ok {
			break
		}
		rv = append(rv, v)
	}
	return rv[:len(rv):len(rv)]
}

func (it *listIterator[E]) MapIndex(f func(int, E) E) ListIterator[E] {
	gen := it.generator
	pos := 0
	return &listIterator[E]{
		generator: func() (E, bool) {
			v, ok := gen()
			if !ok {
				return v, false
			}
			result := f(pos, v)
			pos++
			return result, true
		},
		maxElements: it.maxElements,
	}
}

func (it *listIterator[E]) FilterIndex(pred func(int, E) bool) ListIterator[E] {
	gen := it.generator
	pos := 0
	return &listIterator[E]{
		generator: func() (E, bool) {
			for {
				v, ok := gen()
				if !ok {
					return v, false
				}
				idx := pos
				pos++
				if pred(idx, v) {
					return v, true
				}
			}
		},
		maxElements: it.maxElements,
	}
}

func (it *listIterator[E]) Unique(fn func(E) string) ListIterator[E] {
	gen := it.generator
	seen := make(map[string]struct{}, it.maxElements)
	return &listIterator[E]{
		generator: func() (E, bool) {
			for {
				v, ok := gen()
				if !ok {
					return v, false
				}
				key := fn(v)
				if _, exists := seen[key]; exists {
					continue
				}
				seen[key] = struct{}{}
				return v, true
			}
		},
		maxElements: it.maxElements,
	}
}

func (it *listIterator[E]) ApplyIndex(f func(int, E)) {
	pos := 0
	for {
		v, ok := it.generator()
		if !ok {
			return
		}
		f(pos, v)
		pos++
	}
}

func (it *listIterator[E]) Materialize() ListIterator[E] {
	return FromSlice(it.Collect())
}


1	package collect
2
3	import (
4	"reflect"
5	"slices"
6	)
7
8	// ListGenerator produces values one at a time. It returns (value, true) for
9	// each element and (zero, false) when the sequence is exhausted.
10	type ListGenerator[E any] func() (E, bool)
11
12	// ListIterator is a lazy, pull-based iterator over a sequence of elements.
13	// All intermediate operations (Map, Filter, Sort, Skip, Limit) return a new
14	// iterator; terminal operations (Count, Apply, Collect) consume it.
15	type ListIterator[E any] interface {
16	// Map returns an iterator that applies f to each element.
17	Map(f func(E) E) ListIterator[E]
18	// Filter returns an iterator that yields only elements satisfying pred.
19	Filter(pred func(E) bool) ListIterator[E]
20	// Sort collects all elements, sorts them with less, and returns a new iterator.
21	Sort(less func(E, E) bool) ListIterator[E]
22	// Skip returns an iterator that discards the first count elements.
23	// Panics if count is negative.
24	Skip(count int) ListIterator[E]
25	// Limit returns an iterator that yields at most count elements.
26	// Panics if count is negative.
27	Limit(count int) ListIterator[E]
28	// Count exhausts the iterator and returns the number of elements.
29	Count() int
30	// Apply exhausts the iterator calling f for each element and returns the count.
31	Apply(f func(E)) int
32	// Collect collects all remaining elements into a slice.
33	Collect() []E
34	// MapIndex returns an iterator that applies f with the 0-based index to each element.
35	MapIndex(f func(int, E) E) ListIterator[E]
36	// FilterIndex returns an iterator that yields only elements where pred(index, elem) is true.
37	FilterIndex(pred func(int, E) bool) ListIterator[E]
38	// Unique returns an iterator that deduplicates elements using the string key returned by fn.
39	Unique(fn func(E) string) ListIterator[E]
40	// ApplyIndex exhausts the iterator calling f with the 0-based index for each element.
41	ApplyIndex(f func(int, E))
42	// Materialize materializes the iterator and returns a new iterator over the result.
43	Materialize() ListIterator[E]
44	}
45
46	type listIterator[E any] struct {
47	generator ListGenerator[E]
48	maxElements int
49	}
50
51	// NewListIterator creates a [ListIterator] backed by gen. maxElements is a
52	// hint for the expected upper bound of elements (used for pre-allocation);
53	// pass -1 if unknown.
54	func NewListIterator[E any](gen ListGenerator[E], maxElements int) ListIterator[E] {
55	return &listIterator[E]{generator: gen, maxElements: maxElements}
56	}
57
58	// FromSlice creates a [ListIterator] that yields the elements of inp
59	// in order.
60	func FromSlice[E any](inp []E) ListIterator[E] {
61	pos := 0
62	gen := func() (E, bool) {
63	if pos >= len(inp) {
64	var zero E
65	return zero, false
66	}
67	v := inp[pos]
68	pos++
69	return v, true
70	}
71	return &listIterator[E]{generator: gen, maxElements: len(inp)}
72	}
73
74	// FromSliceWith creates a [ListIterator] that maps each element of inp through cb.
75	func FromSliceWith[E any, I any](inp []E, cb func(E) I) ListIterator[I] {
76	pos := 0
77	gen := func() (I, bool) {
78	if pos >= len(inp) {
79	var zero I
80	return zero, false
81	}
82	v := cb(inp[pos])
83	pos++
84	return v, true
85	}
86	return &listIterator[I]{generator: gen, maxElements: len(inp)}
87	}
88
89	// FromMapKeys creates a [ListIterator] over the keys of a map.
90	func FromMapKeys[K comparable, V any](inp map[K]V) ListIterator[K] {
91	rv := reflect.ValueOf(inp)
92	keys := rv.MapKeys()
93	pos := 0
94	gen := func() (K, bool) {
95	if pos >= len(keys) {
96	var zero K
97	return zero, false
98	}
99	v := keys[pos].Interface().(K)
100	pos++
101	return v, true
102	}
103	return &listIterator[K]{generator: gen, maxElements: len(keys)}
104	}
105
106	// FromMapValues creates a [ListIterator] over the values of a map.
107	func FromMapValues[K comparable, V any](inp map[K]V) ListIterator[V] {
108	rv := reflect.ValueOf(inp)
109	keys := rv.MapKeys()
110	pos := 0
111	gen := func() (V, bool) {
112	if pos >= len(keys) {
113	var zero V
114	return zero, false
115	}
116	v := rv.MapIndex(keys[pos]).Interface().(V)
117	pos++
118	return v, true
119	}
120	return &listIterator[V]{generator: gen, maxElements: len(keys)}
121	}
122
123	func (it *listIterator[E]) Map(f func(E) E) ListIterator[E] {
124	gen := it.generator
125	return &listIterator[E]{
126	generator: func() (E, bool) {
127	v, ok := gen()
128	if !ok {
129	return v, false
130	}
131	return f(v), true
132	},
133	maxElements: it.maxElements,
134	}
135	}
136
137	func (it *listIterator[E]) Filter(pred func(E) bool) ListIterator[E] {
138	gen := it.generator
139	return &listIterator[E]{
140	generator: func() (E, bool) {
141	for {
142	v, ok := gen()
143	if !ok {
144	return v, false
145	}
146	if pred(v) {
147	return v, true
148	}
149	}
150	},
151	maxElements: it.maxElements,
152	}
153	}
154
155	func (it *listIterator[E]) Sort(less func(E, E) bool) ListIterator[E] {
156	all := it.Collect()
157	slices.SortFunc(all, func(a, b E) int {
158	switch {
159	case less(a, b):
160	return -1
161	case less(b, a):
162	return 1
163	default:
164	return 0
165	}
166	})
167	return FromSlice(all)
168	}
169
170	func (it *listIterator[E]) Skip(count int) ListIterator[E] {
171	if count < 0 {
172	panic("collect: Skip count must be non-negative")
173	}
174	gen := it.generator
175	skipped := false
176	return &listIterator[E]{
177	generator: func() (E, bool) {
178	if !skipped {
179	skipped = true
180	for range count {
181	if _, ok := gen(); !ok {
182	var zero E
183	return zero, false
184	}
185	}
186	}
187	return gen()
188	},
189	maxElements: max(it.maxElements-count, 0),
190	}
191	}
192
193	func (it *listIterator[E]) Limit(count int) ListIterator[E] {
194	if count < 0 {
195	panic("collect: Limit count must be non-negative")
196	}
197	gen := it.generator
198	remaining := count
199	return &listIterator[E]{
200	generator: func() (E, bool) {
201	if remaining <= 0 {
202	var zero E
203	return zero, false
204	}
205	remaining--
206	return gen()
207	},
208	maxElements: min(it.maxElements, count),
209	}
210	}
211
212	func (it *listIterator[E]) Count() int {
213	n := 0
214	for {
215	if _, ok := it.generator(); !ok {
216	return n
217	}
218	n++
219	}
220	}
221
222	func (it *listIterator[E]) Apply(f func(E)) int {
223	n := 0
224	for {
225	v, ok := it.generator()
226	if !ok {
227	return n
228	}
229	f(v)
230	n++
231	}
232	}
233
234	func (it *listIterator[E]) Collect() []E {
235	sz := max(it.maxElements, 0)
236	rv := make([]E, 0, sz)
237	for {
238	v, ok := it.generator()
239	if !ok {
240	break
241	}
242	rv = append(rv, v)
243	}
244	return rv[:len(rv):len(rv)]
245	}
246
247	func (it *listIterator[E]) MapIndex(f func(int, E) E) ListIterator[E] {
248	gen := it.generator
249	pos := 0
250	return &listIterator[E]{
251	generator: func() (E, bool) {
252	v, ok := gen()
253	if !ok {
254	return v, false
255	}
256	result := f(pos, v)
257	pos++
258	return result, true
259	},
260	maxElements: it.maxElements,
261	}
262	}
263
264	func (it *listIterator[E]) FilterIndex(pred func(int, E) bool) ListIterator[E] {
265	gen := it.generator
266	pos := 0
267	return &listIterator[E]{
268	generator: func() (E, bool) {
269	for {
270	v, ok := gen()
271	if !ok {
272	return v, false
273	}
274	idx := pos
275	pos++
276	if pred(idx, v) {
277	return v, true
278	}
279	}
280	},
281	maxElements: it.maxElements,
282	}
283	}
284
285	func (it *listIterator[E]) Unique(fn func(E) string) ListIterator[E] {
286	gen := it.generator
287	seen := make(map[string]struct{}, it.maxElements)
288	return &listIterator[E]{
289	generator: func() (E, bool) {
290	for {
291	v, ok := gen()
292	if !ok {
293	return v, false
294	}
295	key := fn(v)
296	if _, exists := seen[key]; exists {
297	continue
298	}
299	seen[key] = struct{}{}
300	return v, true
301	}
302	},
303	maxElements: it.maxElements,
304	}
305	}
306
307	func (it *listIterator[E]) ApplyIndex(f func(int, E)) {
308	pos := 0
309	for {
310	v, ok := it.generator()
311	if !ok {
312	return
313	}
314	f(pos, v)
315	pos++
316	}
317	}
318
319	func (it *listIterator[E]) Materialize() ListIterator[E] {
320	return FromSlice(it.Collect())
321	}
322

iterator.go

Source Files