map_iterator.go

v0.1.0

package collect

import (
	"reflect"
	"slices"
)

// MapGenerator produces key-value pairs one at a time. It returns (key, value, true)
// for each pair and (zeroK, zeroV, false) when the sequence is exhausted.
type MapGenerator[K comparable, V any] func() (K, V, bool)

// MapIterator is a lazy, pull-based iterator over key-value pairs.
// All intermediate operations (MapValues, MapEntries, Filter, Skip, Limit, etc.) return
// a new iterator; terminal operations (Apply, Collect) consume it.
type MapIterator[K comparable, V any] interface {
	// MapValues returns an iterator that transforms values by applying f, keeping keys unchanged.
	MapValues(f func(K, V) V) MapIterator[K, V]
	// MapEntries returns an iterator that transforms both keys and values by applying f.
	MapEntries(f func(K, V) (K, V)) MapIterator[K, V]
	// Filter returns an iterator that yields only pairs satisfying pred.
	Filter(pred func(K, V) bool) MapIterator[K, V]
	// Unique returns an iterator that deduplicates pairs by the string key returned by fn.
	Unique(fn func(K, V) string) MapIterator[K, V]
	// MapValuesIndex returns an iterator that transforms values by applying f with a 0-based index.
	MapValuesIndex(f func(int, K, V) V) MapIterator[K, V]
	// MapEntriesIndex returns an iterator that transforms both keys and values by applying f with a 0-based index.
	MapEntriesIndex(f func(int, K, V) (K, V)) MapIterator[K, V]
	// FilterIndex returns an iterator that yields only pairs satisfying pred, which receives a 0-based index.
	FilterIndex(pred func(int, K, V) bool) MapIterator[K, V]
	// Apply exhausts the iterator, calling f for each key-value pair.
	Apply(f func(K, V))
	// Collect collects all remaining pairs into a map.
	Collect() map[K]V
	// Materialize materializes all pairs into a map and returns a new iterator over it.
	// If less is non-nil the keys are visited in sorted order; otherwise the order is unspecified.
	Materialize(less func(K, K) bool) MapIterator[K, V]
	// Skip returns an iterator that discards the first count pairs.
	// Panics if count is negative.
	Skip(count int) MapIterator[K, V]
	// Limit returns an iterator that yields at most count elements.
	// Panics if count is negative.
	Limit(count int) MapIterator[K, V]
	// Keys returns a ListIterator that yields the keys.
	Keys() ListIterator[K]
	// Values returns a ListIterator that yields the values.
	Values() ListIterator[V]
}

var _ MapIterator[int, any] = &mapIterator[int, any]{}

type mapIterator[K comparable, V any] struct {
	generator   MapGenerator[K, V]
	maxElements int
}

// NewMapIterator creates a [MapIterator] backed by gen. maxElements is a hint
// for the expected upper bound of pairs (used for pre-allocation); pass -1 if unknown.
func NewMapIterator[K comparable, V any](gen MapGenerator[K, V], maxElements int) MapIterator[K, V] {
	return &mapIterator[K, V]{generator: gen, maxElements: maxElements}
}

// FromMap creates a [MapIterator] that yields the entries of inp.
// If less is non-nil the keys are visited in sorted order; otherwise the
// iteration order is unspecified (uses reflect.MapRange).
func FromMap[K comparable, V any](inp map[K]V, less func(K, K) bool) MapIterator[K, V] {
	if less != nil {
		keys := make([]K, 0, len(inp))
		for k := range inp {
			keys = append(keys, k)
		}
		slices.SortFunc(keys, func(a, b K) int {
			switch {
			case less(a, b):
				return -1
			case less(b, a):
				return 1
			default:
				return 0
			}
		})
		pos := 0
		gen := func() (K, V, bool) {
			if pos >= len(keys) {
				var zeroK K
				var zeroV V
				return zeroK, zeroV, false
			}
			k := keys[pos]
			pos++
			return k, inp[k], true
		}
		return &mapIterator[K, V]{generator: gen, maxElements: len(inp)}
	}

	iter := reflect.ValueOf(inp).MapRange()
	gen := func() (K, V, bool) {
		if !iter.Next() {
			var zeroK K
			var zeroV V
			return zeroK, zeroV, false
		}
		return iter.Key().Interface().(K), iter.Value().Interface().(V), true
	}
	return &mapIterator[K, V]{generator: gen, maxElements: len(inp)}
}

func (m *mapIterator[K, V]) MapValues(f func(K, V) V) MapIterator[K, V] {
	gen := m.generator
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			k, v, ok := gen()
			if !ok {
				return k, v, false
			}
			return k, f(k, v), true
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) MapEntries(f func(K, V) (K, V)) MapIterator[K, V] {
	gen := m.generator
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			k, v, ok := gen()
			if !ok {
				return k, v, false
			}
			nk, nv := f(k, v)
			return nk, nv, true
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) Filter(pred func(K, V) bool) MapIterator[K, V] {
	gen := m.generator
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			for {
				k, v, ok := gen()
				if !ok {
					return k, v, false
				}
				if pred(k, v) {
					return k, v, true
				}
			}
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) Unique(fn func(K, V) string) MapIterator[K, V] {
	gen := m.generator
	sz := m.maxElements
	if sz < 0 {
		sz = 0
	}
	seen := make(map[string]struct{}, sz)
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			for {
				k, v, ok := gen()
				if !ok {
					return k, v, false
				}
				key := fn(k, v)
				if _, exists := seen[key]; exists {
					continue
				}
				seen[key] = struct{}{}
				return k, v, true
			}
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) MapValuesIndex(f func(int, K, V) V) MapIterator[K, V] {
	gen := m.generator
	pos := 0
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			k, v, ok := gen()
			if !ok {
				return k, v, false
			}
			nv := f(pos, k, v)
			pos++
			return k, nv, true
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) MapEntriesIndex(f func(int, K, V) (K, V)) MapIterator[K, V] {
	gen := m.generator
	pos := 0
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			k, v, ok := gen()
			if !ok {
				return k, v, false
			}
			nk, nv := f(pos, k, v)
			pos++
			return nk, nv, true
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) FilterIndex(pred func(int, K, V) bool) MapIterator[K, V] {
	gen := m.generator
	pos := 0
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			for {
				k, v, ok := gen()
				if !ok {
					return k, v, false
				}
				idx := pos
				pos++
				if pred(idx, k, v) {
					return k, v, true
				}
			}
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) Apply(f func(K, V)) {
	for {
		k, v, ok := m.generator()
		if !ok {
			return
		}
		f(k, v)
	}
}

func (m *mapIterator[K, V]) Collect() map[K]V {
	sz := m.maxElements
	if sz < 0 {
		sz = 0
	}
	rv := make(map[K]V, sz)
	for {
		k, v, ok := m.generator()
		if !ok {
			break
		}
		rv[k] = v
	}
	return rv
}

func (m *mapIterator[K, V]) Materialize(less func(K, K) bool) MapIterator[K, V] {
	return FromMap(m.Collect(), less)
}

func (m *mapIterator[K, V]) Skip(count int) MapIterator[K, V] {
	if count < 0 {
		panic("collect: Skip count must be non-negative")
	}
	gen := m.generator
	skipped := false
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			if !skipped {
				skipped = true
				for range count {
					if _, _, ok := gen(); !ok {
						var zeroK K
						var zeroV V
						return zeroK, zeroV, false
					}
				}
			}
			return gen()
		},
		maxElements: max(m.maxElements-count, 0),
	}
}

func (m *mapIterator[K, V]) Limit(count int) MapIterator[K, V] {
	if count < 0 {
		panic("collect: Limit count must be non-negative")
	}
	gen := m.generator
	remaining := count
	return &mapIterator[K, V]{
		generator: func() (K, V, bool) {
			if remaining <= 0 {
				var zeroK K
				var zeroV V
				return zeroK, zeroV, false
			}
			remaining--
			return gen()
		},
		maxElements: min(m.maxElements, count),
	}
}

func (m *mapIterator[K, V]) Keys() ListIterator[K] {
	gen := m.generator
	return &listIterator[K]{
		generator: func() (K, bool) {
			k, _, ok := gen()
			return k, ok
		},
		maxElements: m.maxElements,
	}
}

func (m *mapIterator[K, V]) Values() ListIterator[V] {
	gen := m.generator
	return &listIterator[V]{
		generator: func() (V, bool) {
			_, v, ok := gen()
			return v, ok
		},
		maxElements: m.maxElements,
	}
}


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

map_iterator.go

Source Files