internal/iterutil: add Chain, Chain2 iterators

Add utilities to concatenate multiple iterators of the same type into a
single iterator.

Signed-off-by: Cory Snider <csnider@mirantis.com>
This commit is contained in:
Cory Snider
2025-09-05 14:52:09 -04:00
parent 2e12287747
commit 65ffac3dbf
2 changed files with 66 additions and 0 deletions

View File

@@ -27,3 +27,29 @@ func Deref[T any, P *T](s iter.Seq[P]) iter.Seq[T] {
}
}
}
// Chain concatenates multiple iterators into a single iterator.
func Chain[T any](iters ...iter.Seq[T]) iter.Seq[T] {
return func(yield func(T) bool) {
for _, it := range iters {
for v := range it {
if !yield(v) {
return
}
}
}
}
}
// Chain2 concatenates multiple iterators into a single iterator.
func Chain2[K, V any](iters ...iter.Seq2[K, V]) iter.Seq2[K, V] {
return func(yield func(K, V) bool) {
for _, it := range iters {
for k, v := range it {
if !yield(k, v) {
return
}
}
}
}
}

View File

@@ -1,6 +1,7 @@
package iterutil
import (
"maps"
"slices"
"testing"
@@ -26,3 +27,42 @@ func TestDeref(t *testing.T) {
b := slices.Collect(Deref(slices.Values(a)))
assert.DeepEqual(t, b, []int{0, 1, 2})
}
func TestChain(t *testing.T) {
a := []int{1, 2, 3}
b := []int{4, 5}
c := []int{6}
ab := Chain(slices.Values(a), slices.Values(b))
abc := Chain(ab, slices.Values(c))
assert.DeepEqual(t, slices.Collect(ab), []int{1, 2, 3, 4, 5})
assert.DeepEqual(t, slices.Collect(abc), []int{1, 2, 3, 4, 5, 6})
}
func TestChain2(t *testing.T) {
a := map[string]int{
"a": 1,
"b": 2,
"c": 3,
}
b := map[string]int{
"d": 4,
"e": 5,
}
c := map[string]int{
"f": 6,
}
ab := Chain2(maps.All(a), maps.All(b))
abc := Chain2(ab, maps.All(c))
expab := maps.Clone(a)
maps.Insert(expab, maps.All(b))
expabc := maps.Clone(expab)
maps.Insert(expabc, maps.All(c))
assert.DeepEqual(t, maps.Collect(ab), expab)
assert.DeepEqual(t, maps.Collect(abc), expabc)
}