From b9f085495091c567eb8e81c1b07633f9a540b598 Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Mon, 26 Dec 2022 15:58:44 +0800
Subject: [PATCH] feat: add join iterator

---
 iterator/operation.go      | 42 ++++++++++++++++++++++++++++++++++++++
 iterator/operation_test.go | 15 ++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/iterator/operation.go b/iterator/operation.go
index b77389c..f0f2e65 100644
--- a/iterator/operation.go
+++ b/iterator/operation.go
@@ -59,3 +59,45 @@ func (fr *filterIterator[T]) Next() (T, bool) {
 func (fr *filterIterator[T]) HasNext() bool {
 	return fr.iter.HasNext()
 }
+
+// Join creates an iterator that join all elements of iters[0], then all elements of iters[1] and so on.
+func Join[T any](iters ...Iterator[T]) Iterator[T] {
+	return &joinIterator[T]{
+		iters: iters,
+	}
+}
+
+type joinIterator[T any] struct {
+	iters []Iterator[T]
+}
+
+func (iter *joinIterator[T]) Next() (T, bool) {
+	for len(iter.iters) > 0 {
+		item, ok := iter.iters[0].Next()
+		if ok {
+			return item, true
+		}
+		iter.iters = iter.iters[1:]
+	}
+	var zero T
+	return zero, false
+}
+
+func (iter *joinIterator[T]) HasNext() bool {
+	if len(iter.iters) == 0 {
+		return false
+	}
+	if len(iter.iters) == 1 {
+		return iter.iters[0].HasNext()
+	}
+
+	result := iter.iters[0].HasNext()
+
+	for i := 1; i < len(iter.iters); i++ {
+		it := iter.iters[i]
+		hasNext := it.HasNext()
+		result = result || hasNext
+	}
+
+	return result
+}
diff --git a/iterator/operation_test.go b/iterator/operation_test.go
index 582e23d..89cd529 100644
--- a/iterator/operation_test.go
+++ b/iterator/operation_test.go
@@ -37,3 +37,18 @@ func TestFilterIterator(t *testing.T) {
 	result := ToSlice(iter)
 	assert.Equal([]int{1, 2}, result)
 }
+
+func TestJoinIterator(t *testing.T) {
+	assert := internal.NewAssert(t, "TestJoinIterator")
+
+	iter1 := FromSlice([]int{1, 2})
+	iter2 := FromSlice([]int{3, 4})
+
+	iter := Join(iter1, iter2)
+
+	item, ok := iter.Next()
+	assert.Equal(1, item)
+	assert.Equal(true, ok)
+
+	assert.Equal([]int{2, 3, 4}, ToSlice(iter))
+}