From 65315dafb193b2095943046b48a5f07394edb50f Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Mon, 26 Dec 2022 17:20:14 +0800
Subject: [PATCH] feat: add take iterator

---
 iterator/operation.go      | 25 +++++++++++++++++++++++++
 iterator/operation_test.go | 11 +++++++++++
 2 files changed, 36 insertions(+)

diff --git a/iterator/operation.go b/iterator/operation.go
index a2ec362..6340c5a 100644
--- a/iterator/operation.go
+++ b/iterator/operation.go
@@ -112,3 +112,28 @@ func Reduce[T any, U any](iter Iterator[T], initial U, reducer func(U, T) U) U {
 
 	return acc
 }
+
+func Take[T any](it Iterator[T], num int) Iterator[T] {
+	return &takeIterator[T]{it: it, num: num}
+}
+
+type takeIterator[T any] struct {
+	it  Iterator[T]
+	num int
+}
+
+func (iter *takeIterator[T]) Next() (T, bool) {
+	if iter.num <= 0 {
+		var zero T
+		return zero, false
+	}
+	item, ok := iter.it.Next()
+	if ok {
+		iter.num--
+	}
+	return item, ok
+}
+
+func (iter *takeIterator[T]) HasNext() bool {
+	return iter.num > 0
+}
diff --git a/iterator/operation_test.go b/iterator/operation_test.go
index 71e184a..b9052aa 100644
--- a/iterator/operation_test.go
+++ b/iterator/operation_test.go
@@ -60,3 +60,14 @@ func TestReduce(t *testing.T) {
 	sum := Reduce(iter, 0, func(a, b int) int { return a + b })
 	assert.Equal(10, sum)
 }
+
+func TestTakeIterator(t *testing.T) {
+	assert := internal.NewAssert(t, "TestTakeIterator")
+
+	iter := FromSlice([]int{1, 2, 3, 4, 5})
+
+	iter = Take(iter, 3)
+
+	result := ToSlice(iter)
+	assert.Equal([]int{1, 2, 3}, result)
+}