From c104272a33e4db03fa68d8cb6a93d015f650fb42 Mon Sep 17 00:00:00 2001
From: x5a17ed <0x5a17ed@tuta.io>
Date: Fri, 24 May 2024 17:34:35 +0200
Subject: [PATCH] add chunk iterator

---
 itlib/chunk.go      | 57 ++++++++++++++++++++++++++
 itlib/chunk_test.go | 98 +++++++++++++++++++++++++++++++++++++++++++++
 itlib/limit.go      |  6 ++-
 3 files changed, 160 insertions(+), 1 deletion(-)
 create mode 100644 itlib/chunk.go
 create mode 100644 itlib/chunk_test.go

diff --git a/itlib/chunk.go b/itlib/chunk.go
new file mode 100644
index 0000000..b229960
--- /dev/null
+++ b/itlib/chunk.go
@@ -0,0 +1,57 @@
+// Copyright (c) 2024 individual contributors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// <https://www.apache.org/licenses/LICENSE-2.0>
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package itlib
+
+import (
+	"github.com/0x5a17ed/itkit"
+)
+
+// ChunkIterator yields iterators yielding up to n items from a source
+// iterator until the source iterator is exhausted.
+type ChunkIterator[T any] struct {
+	src *PeekIterator[T]
+
+	cur *LimitIterator[T]
+	n   uint
+}
+
+// Ensure ChunkIterator implements the iterator interface.
+var _ itkit.Iterator[itkit.Iterator[struct{}]] = &ChunkIterator[struct{}]{}
+
+// Next implements the [itkit.Iterator.Next] interface.
+func (it *ChunkIterator[T]) Next() bool {
+	if it.cur != nil && it.cur.n > 0 {
+		// Drop any items not consumed in the previous chunk
+		// to make sure that the next chunk only contains
+		// new items.
+		Drop(it.cur.n, it.src.Iter())
+	}
+
+	if _, ok := it.src.Peek(); !ok {
+		return false
+	}
+	it.cur = newLimitIterator(it.n, it.src.Iter())
+	return true
+}
+
+// Value implements the [itkit.Iterator.Value] interface.
+func (it *ChunkIterator[T]) Value() itkit.Iterator[T] {
+	return it.cur
+}
+
+// Chunk returns a new [ChunkIterator] value.
+func Chunk[T any](n uint, src itkit.Iterator[T]) itkit.Iterator[itkit.Iterator[T]] {
+	return &ChunkIterator[T]{src: newPeekIterator(src), n: n}
+}
diff --git a/itlib/chunk_test.go b/itlib/chunk_test.go
new file mode 100644
index 0000000..602b311
--- /dev/null
+++ b/itlib/chunk_test.go
@@ -0,0 +1,98 @@
+// Copyright (c) 2024 individual contributors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// <https://www.apache.org/licenses/LICENSE-2.0>
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package itlib_test
+
+import (
+	"testing"
+
+	"github.com/0x5a17ed/itkit"
+	"github.com/0x5a17ed/itkit/iters/sliceit"
+	"github.com/0x5a17ed/itkit/itlib"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestChunk(t *testing.T) {
+	tt := []struct {
+		name string
+		inp  []string
+		tx   func(el itkit.Iterator[string]) itkit.Iterator[string]
+		exp  [][]string
+	}{
+		{
+			name: "emtpy",
+			inp:  nil,
+			exp:  nil,
+		},
+
+		{
+			name: "partial-first",
+			inp:  []string{"A", "B"},
+			exp:  [][]string{{"A", "B"}},
+		},
+
+		{
+			name: "complete",
+			inp: []string{
+				"A", "B", "C", "D", "E", "F",
+			},
+			exp: [][]string{
+				{"A", "B", "C"},
+				{"D", "E", "F"},
+			},
+		},
+
+		{
+			name: "partial-last",
+			inp: []string{
+				"A", "B", "C", "D", "E", "F", "G",
+			},
+			exp: [][]string{
+				{"A", "B", "C"},
+				{"D", "E", "F"},
+				{"G"},
+			},
+		},
+
+		{
+			name: "half-consumed",
+			inp: []string{
+				"A", "B", "C", "D", "E", "F", "G",
+			},
+			tx: func(el itkit.Iterator[string]) itkit.Iterator[string] {
+				return itlib.Limit(el, 2)
+			},
+			exp: [][]string{
+				{"A", "B"},
+				{"D", "E"},
+				{"G"},
+			},
+		},
+	}
+
+	for _, tc := range tt {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			it := itlib.Chunk(3, sliceit.In(tc.inp))
+
+			if tc.tx != nil {
+				it = itlib.Map(it, tc.tx)
+			}
+
+			got := sliceit.To(itlib.Map(it, sliceit.To[string]))
+
+			assert.Equal(t, tc.exp, got)
+		})
+	}
+}
diff --git a/itlib/limit.go b/itlib/limit.go
index 12a6e79..1af1622 100644
--- a/itlib/limit.go
+++ b/itlib/limit.go
@@ -42,7 +42,11 @@ func (it *LimitIterator[T]) Value() (v T) {
 	return it.src.Value()
 }
 
+func newLimitIterator[T any](n uint, src itkit.Iterator[T]) *LimitIterator[T] {
+	return &LimitIterator[T]{n: n, src: src}
+}
+
 // Limit returns a new [LimitIterator] instance.
 func Limit[T any](src itkit.Iterator[T], n uint) itkit.Iterator[T] {
-	return &LimitIterator[T]{src: src, n: n}
+	return newLimitIterator(n, src)
 }