From 65ffac3dbff5d2314994259d48fe702cc80fcc46 Mon Sep 17 00:00:00 2001
From: Cory Snider <csnider@mirantis.com>
Date: Fri, 5 Sep 2025 14:52:09 -0400
Subject: [PATCH] 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>
---
 internal/iterutil/iterutil.go      | 26 +++++++++++++++++++
 internal/iterutil/iterutil_test.go | 40 ++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/internal/iterutil/iterutil.go b/internal/iterutil/iterutil.go
index 2a29e0f437..5a5296f779 100644
--- a/internal/iterutil/iterutil.go
+++ b/internal/iterutil/iterutil.go
@@ -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
+				}
+			}
+		}
+	}
+}
diff --git a/internal/iterutil/iterutil_test.go b/internal/iterutil/iterutil_test.go
index 4753108b21..c520d5a8c7 100644
--- a/internal/iterutil/iterutil_test.go
+++ b/internal/iterutil/iterutil_test.go
@@ -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)
+}