map_iterator.go

v0.6.1

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 := max(m.maxElements, 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 := max(m.maxElements, 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 := max(m.maxElements, 0)
157	seen := make(map[string]struct{}, sz)
158	return &mapIterator[K, V]{
159	generator: func() (K, V, bool) {
160	for {
161	k, v, ok := gen()
162	if !ok {
163	return k, v, false
164	}
165	key := fn(k, v)
166	if _, exists := seen[key]; exists {
167	continue
168	}
169	seen[key] = struct{}{}
170	return k, v, true
171	}
172	},
173	maxElements: m.maxElements,
174	}
175	}
176
177	func (m *mapIterator[K, V]) MapValuesIndex(f func(int, K, V) V) MapIterator[K, V] {
178	gen := m.generator
179	pos := 0
180	return &mapIterator[K, V]{
181	generator: func() (K, V, bool) {
182	k, v, ok := gen()
183	if !ok {
184	return k, v, false
185	}
186	nv := f(pos, k, v)
187	pos++
188	return k, nv, true
189	},
190	maxElements: m.maxElements,
191	}
192	}
193
194	func (m *mapIterator[K, V]) MapEntriesIndex(f func(int, K, V) (K, V)) MapIterator[K, V] {
195	gen := m.generator
196	pos := 0
197	return &mapIterator[K, V]{
198	generator: func() (K, V, bool) {
199	k, v, ok := gen()
200	if !ok {
201	return k, v, false
202	}
203	nk, nv := f(pos, k, v)
204	pos++
205	return nk, nv, true
206	},
207	maxElements: m.maxElements,
208	}
209	}
210
211	func (m *mapIterator[K, V]) FilterIndex(pred func(int, K, V) bool) MapIterator[K, V] {
212	gen := m.generator
213	pos := 0
214	return &mapIterator[K, V]{
215	generator: func() (K, V, bool) {
216	for {
217	k, v, ok := gen()
218	if !ok {
219	return k, v, false
220	}
221	idx := pos
222	pos++
223	if pred(idx, k, v) {
224	return k, v, true
225	}
226	}
227	},
228	maxElements: m.maxElements,
229	}
230	}
231
232	func (m *mapIterator[K, V]) Apply(f func(K, V)) {
233	for {
234	k, v, ok := m.generator()
235	if !ok {
236	return
237	}
238	f(k, v)
239	}
240	}
241
242	func (m *mapIterator[K, V]) Collect() map[K]V {
243	sz := max(m.maxElements, 0)
244	rv := make(map[K]V, sz)
245	for {
246	k, v, ok := m.generator()
247	if !ok {
248	break
249	}
250	rv[k] = v
251	}
252	return rv
253	}
254
255	func (m *mapIterator[K, V]) Materialize(less func(K, K) bool) MapIterator[K, V] {
256	return FromMap(m.Collect(), less)
257	}
258
259	func (m *mapIterator[K, V]) Skip(count int) MapIterator[K, V] {
260	if count < 0 {
261	panic("collect: Skip count must be non-negative")
262	}
263	gen := m.generator
264	skipped := false
265	return &mapIterator[K, V]{
266	generator: func() (K, V, bool) {
267	if !skipped {
268	skipped = true
269	for range count {
270	if _, _, ok := gen(); !ok {
271	var zeroK K
272	var zeroV V
273	return zeroK, zeroV, false
274	}
275	}
276	}
277	return gen()
278	},
279	maxElements: max(m.maxElements-count, 0),
280	}
281	}
282
283	func (m *mapIterator[K, V]) Limit(count int) MapIterator[K, V] {
284	if count < 0 {
285	panic("collect: Limit count must be non-negative")
286	}
287	gen := m.generator
288	remaining := count
289	return &mapIterator[K, V]{
290	generator: func() (K, V, bool) {
291	if remaining <= 0 {
292	var zeroK K
293	var zeroV V
294	return zeroK, zeroV, false
295	}
296	remaining--
297	return gen()
298	},
299	maxElements: min(m.maxElements, count),
300	}
301	}
302
303	func (m *mapIterator[K, V]) Keys() ListIterator[K] {
304	gen := m.generator
305	return &listIterator[K]{
306	generator: func() (K, bool) {
307	k, _, ok := gen()
308	return k, ok
309	},
310	maxElements: m.maxElements,
311	}
312	}
313
314	func (m *mapIterator[K, V]) Values() ListIterator[V] {
315	gen := m.generator
316	return &listIterator[V]{
317	generator: func() (V, bool) {
318	_, v, ok := gen()
319	return v, ok
320	},
321	maxElements: m.maxElements,
322	}
323	}
324

map_iterator.go

Source Files