test: add comprehensive security tests for leaks and panics

README.md

@@ -55,6 +55,72 @@ func main() {
 
 ---
 
+## **Best Practices & Important Notes**
+
+### **⚠️ Always Close Iterators When Using Next/Current**
+
+When using imperative-style iteration with `Next()` and `Current()`, **always call `Close()`** or use `defer iter.Close()` to prevent goroutine leaks:
+
+```go
+// ✅ CORRECT: Use defer to ensure cleanup
+iter := itertools.Range(0, 1000000)
+defer iter.Close()  // Always clean up!
+
+for iter.Next() {
+    if someCondition {
+        break  // Close() will be called via defer
+    }
+    process(iter.Current())
+}
+```
+
+```go
+// ❌ WRONG: Goroutine leak if you break early
+iter := itertools.Range(0, 1000000)
+for iter.Next() {
+    if someCondition {
+        break  // LEAK! Goroutine continues running
+    }
+}
+```
+
+**Note**: Functional-style operations (`Collect()`, `Filter()`, `Map()`, etc.) handle cleanup automatically.
+
+### **⚠️ Avoid Unbounded Operations on Infinite Iterators**
+
+Some operations consume the entire sequence and will cause out-of-memory errors with infinite iterators:
+
+**Unbounded operations**:
+- `Collect()` - Loads all elements into a slice
+- `Unique()` - Builds an unbounded map
+- `GroupBy()` - Builds unbounded maps and slices
+- `Union()`, `Intersection()`, `Difference()` - Build unbounded maps
+- `Partition()`, `Reverse()`, `Sort()`, `Shuffle()` - Collect all elements
+- `Cycle()` - Creates infinite repetition
+
+```go
+// ❌ WRONG: Out of memory!
+infiniteIter := itertools.FromFunc(func() (int, bool) {
+    return rand.Int(), true
+})
+infiniteIter.Collect()  // OOM!
+
+// ✅ CORRECT: Use Take to limit elements
+infiniteIter := itertools.FromFunc(func() (int, bool) {
+    return rand.Int(), true
+})
+result := infiniteIter.Take(1000).Collect()  // OK
+```
+
+### **Input Validation**
+
+Methods validate inputs and panic on invalid arguments:
+- `StepBy(n)` - Panics if `n <= 0`
+- `ChunkSlice(size)`, `Chunks(size)`, `ChunkList(size)` - Panic if `size <= 0`
+- `Take(n)`, `Drop(n)`, `Nth(n)` - Negative values are treated as 0
+
+---
+
 ## **Features**
 - **Chainable API**: Combine transformations like `Filter`, `Map`, `Take`, and `Drop` into one functional-style chain.
 - **Laziness**: Iterators are lazy; they only compute elements as needed.

iterator.go

@@ -119,6 +119,36 @@ func (it *Iterator[V]) Current() V {
 	return *it.curr
 }
 
+// Close releases any resources held by the iterator, particularly cleaning up
+// the goroutine spawned by iter.Pull() when using Next/Current iteration.
+//
+// IMPORTANT: Always call Close() when using Next/Current iteration and not
+// fully exhausting the iterator. Failing to do so will leak a goroutine.
+//
+// Close is idempotent - it can be called multiple times safely.
+// Close is automatically called when the iterator is fully exhausted.
+//
+// Best Practice - Use defer for guaranteed cleanup:
+//
+//	iter := itertools.Range(0, 1000000)
+//	defer iter.Close()
+//	for iter.Next() {
+//	    if someCondition {
+//	        break // Close() will be called via defer
+//	    }
+//	    process(iter.Current())
+//	}
+//
+// Note: Close only affects Next/Current iteration. Functional-style operations
+// (Collect, Each, Filter, Map, etc.) handle cleanup automatically.
+func (it *Iterator[V]) Close() {
+	if it.stop != nil {
+		it.stop()
+		it.stop = nil // Make Close idempotent
+	}
+	it.done = true
+}
+
 // Collect consumes the iterator and returns all elements as a slice.
 // After calling Collect, the iterator is exhausted.
 //
@@ -251,6 +281,8 @@ func (it *Iterator[V]) Chain(other *Iterator[V]) *Iterator[V] {
 // Take returns a new iterator that yields at most the first n elements.
 // If the iterator has fewer than n elements, all elements are yielded.
 //
+// If n is negative, it is treated as 0 (returns an empty iterator).
+//
 // Take is lazy and supports early termination.
 //
 // Example:
@@ -259,6 +291,9 @@ func (it *Iterator[V]) Chain(other *Iterator[V]) *Iterator[V] {
 //	first3 := iter.Take(3).Collect()
 //	// first3 is []int{1, 2, 3}
 func (it *Iterator[V]) Take(n int) *Iterator[V] {
+	if n < 0 {
+		n = 0
+	}
 	return &Iterator[V]{
 		seq: func(yield func(V) bool) {
 			i := 0
@@ -276,12 +311,17 @@ func (it *Iterator[V]) Take(n int) *Iterator[V] {
 // Drop returns a new iterator that skips the first n elements and yields the rest.
 // If the iterator has n or fewer elements, the resulting iterator is empty.
 //
+// If n is negative, it is treated as 0 (yields all elements).
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{1, 2, 3, 4, 5})
 //	rest := iter.Drop(2).Collect()
 //	// rest is []int{3, 4, 5}
 func (it *Iterator[V]) Drop(n int) *Iterator[V] {
+	if n < 0 {
+		n = 0
+	}
 	return &Iterator[V]{
 		seq: func(yield func(V) bool) {
 			i := 0
@@ -375,13 +415,18 @@ func (it *Iterator[V]) Last() V {
 // Nth returns the nth element (0-indexed) of the iterator.
 // It panics if the iterator has fewer than n+1 elements.
 //
+// If n is negative, it is treated as 0 (returns the first element).
+//
 // For a safe alternative that doesn't panic, use NthOr.
 //
 // Example:
 //
 //	iter := itertools.ToIter([]int{10, 20, 30, 40})
 //	third := iter.Nth(2) // Returns 30
 func (it *Iterator[V]) Nth(n int) V {
+	if n < 0 {
+		n = 0
+	}
 	return it.Drop(n).First()
 }
 
@@ -438,12 +483,17 @@ func (it *Iterator[V]) LastOr(defaultValue V) V {
 // NthOr returns the nth element (0-indexed) of the iterator, or defaultValue if there aren't enough elements.
 // This is a safe alternative to Nth that doesn't panic.
 //
+// If n is negative, it is treated as 0 (returns the first element, or defaultValue if empty).
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{10, 20, 30})
 //	third := iter.NthOr(2, 999)  // Returns 30
 //	fifth := iter.NthOr(4, 999)  // Returns 999 (not enough elements)
 func (it *Iterator[V]) NthOr(n int, defaultValue V) V {
+	if n < 0 {
+		n = 0
+	}
 	var result V
 	found := false
 	i := 0
@@ -802,12 +852,17 @@ func (it *Iterator[V]) Difference(other *Iterator[V], keyFunc func(V) any) *Iter
 
 // StepBy returns an iterator that yields every nth element (0-indexed).
 //
+// Panics if n <= 0.
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
 //	everyThird := iter.StepBy(3).Collect()
 //	// everyThird is []int{0, 3, 6, 9}
 func (it *Iterator[V]) StepBy(n int) *Iterator[V] {
+	if n <= 0 {
+		panic(fmt.Sprintf("StepBy: step must be positive, got %d", n))
+	}
 	return &Iterator[V]{
 		seq: func(yield func(V) bool) {
 			i := 0

itertools.go

@@ -2,6 +2,7 @@ package itertools
 
 import (
 	"cmp"
+	"fmt"
 	"iter"
 
 	"golang.org/x/exp/constraints"
@@ -199,12 +200,17 @@ func Product[V any, T Productable](it *Iterator[V], transform func(V) T, one T)
 //
 // Each chunk is a separate slice - modifications won't affect the original data.
 //
+// Panics if size <= 0.
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{1, 2, 3, 4, 5, 6, 7})
 //	chunks := itertools.ChunkSlice(iter, 3).Collect()
 //	// chunks is [][]int{{1, 2, 3}, {4, 5, 6}, {7}}
 func ChunkSlice[V any](it *Iterator[V], size int) *Iterator[[]V] {
+	if size <= 0 {
+		panic(fmt.Sprintf("ChunkSlice: size must be positive, got %d", size))
+	}
 	return &Iterator[[]V]{
 		seq: func(yield func([]V) bool) {
 			chunk := make([]V, 0, size)
@@ -236,6 +242,8 @@ func ChunkSlice[V any](it *Iterator[V], size int) *Iterator[[]V] {
 //
 // Unlike ChunkSlice, this returns iterators instead of slices.
 //
+// Panics if size <= 0.
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{1, 2, 3, 4, 5})
@@ -248,6 +256,9 @@ func ChunkSlice[V any](it *Iterator[V], size int) *Iterator[[]V] {
 //	//         [3 4]
 //	//         [5]
 func Chunks[V any](it *Iterator[V], size int) *Iterator[*Iterator[V]] {
+	if size <= 0 {
+		panic(fmt.Sprintf("Chunks: size must be positive, got %d", size))
+	}
 	return &Iterator[*Iterator[V]]{
 		seq: func(yield func(*Iterator[V]) bool) {
 			chunk := make([]V, 0, size)
@@ -277,12 +288,17 @@ func Chunks[V any](it *Iterator[V], size int) *Iterator[*Iterator[V]] {
 // ChunkList returns a slice of iterators, each containing up to `size` elements.
 // This is a convenience function that collects Chunks into a slice.
 //
+// Panics if size <= 0.
+//
 // Example:
 //
 //	iter := itertools.ToIter([]int{1, 2, 3, 4, 5, 6})
 //	chunks := itertools.ChunkList(iter, 2)
 //	// chunks is []*Iterator with 3 iterators containing [1,2], [3,4], [5,6]
 func ChunkList[V any](it *Iterator[V], size int) []*Iterator[V] {
+	if size <= 0 {
+		panic(fmt.Sprintf("ChunkList: size must be positive, got %d", size))
+	}
 	return Chunks(it, size).Collect()
 }
 

itertools_test.go

@@ -363,3 +363,98 @@ func TestChunkList_Functional(t *testing.T) {
 	assert.Equal(t, []int{6, 7, 8}, chunks[2].Collect())
 	assert.Equal(t, []int{9}, chunks[3].Collect())
 }
+
+// =============================================================================
+// PANIC AND EDGE CASE TESTS
+// =============================================================================
+
+// TestChunkSlice_PanicOnNegativeSize verifies ChunkSlice panics with negative size
+func TestChunkSlice_PanicOnNegativeSize(t *testing.T) {
+	iter := itertools.Range(0, 10)
+
+	assert.Panics(t, func() {
+		itertools.ChunkSlice(iter, -5).Collect()
+	}, "ChunkSlice should panic with negative size")
+}
+
+// TestChunks_PanicOnNegativeSize verifies Chunks panics with negative size
+func TestChunks_PanicOnNegativeSize(t *testing.T) {
+	iter := itertools.Range(0, 10)
+
+	assert.Panics(t, func() {
+		itertools.Chunks(iter, -3).Collect()
+	}, "Chunks should panic with negative size")
+}
+
+// TestChunkList_PanicOnNegativeSize verifies ChunkList panics with negative size
+func TestChunkList_PanicOnNegativeSize(t *testing.T) {
+	iter := itertools.Range(0, 10)
+
+	assert.Panics(t, func() {
+		itertools.ChunkList(iter, -2)
+	}, "ChunkList should panic with negative size")
+}
+
+// TestChunkSlice_ZeroSize tests that ChunkSlice panics with size=0
+func TestChunkSlice_ZeroSize(t *testing.T) {
+	iter := itertools.Range(0, 10)
+
+	// Now correctly panics with zero size
+	assert.Panics(t, func() {
+		itertools.ChunkSlice(iter, 0).Collect()
+	}, "ChunkSlice should panic with zero size")
+}
+
+// TestChunks_ZeroSize tests that Chunks panics with size=0
+func TestChunks_ZeroSize(t *testing.T) {
+	iter := itertools.Range(0, 10)
+
+	// Now correctly panics with zero size
+	assert.Panics(t, func() {
+		itertools.Chunks(iter, 0).Collect()
+	}, "Chunks should panic with zero size")
+}
+
+// TestFold_BasicOperation tests Fold with basic addition
+func TestFold_BasicOperation(t *testing.T) {
+	iter := itertools.Range(1, 6) // 1, 2, 3, 4, 5
+
+	sum := itertools.Fold(iter, func(acc, v int) int {
+		return acc + v
+	}, 0)
+
+	assert.Equal(t, 15, sum)
+}
+
+// TestFold_StringConcatenation tests Fold with string concatenation
+func TestFold_StringConcatenation(t *testing.T) {
+	iter := itertools.ToIter([]string{"a", "b", "c"})
+
+	result := itertools.Fold(iter, func(acc, v string) string {
+		return acc + v
+	}, "")
+
+	assert.Equal(t, "abc", result)
+}
+
+// TestCartesianProduct_EmptyIterators tests CartesianProduct with empty inputs
+func TestCartesianProduct_EmptyIterators(t *testing.T) {
+	iter1 := itertools.ToIter([]int{})
+	iter2 := itertools.ToIter([]string{"a", "b"})
+
+	result := itertools.CartesianProduct(iter1, iter2).Collect()
+
+	assert.Empty(t, result, "CartesianProduct with empty iterator should return empty")
+}
+
+// TestCartesianProduct_SingleElements tests CartesianProduct with single elements
+func TestCartesianProduct_SingleElements(t *testing.T) {
+	iter1 := itertools.ToIter([]int{1})
+	iter2 := itertools.ToIter([]string{"x"})
+
+	result := itertools.CartesianProduct(iter1, iter2).Collect()
+
+	assert.Equal(t, 1, len(result))
+	assert.Equal(t, 1, result[0].X)
+	assert.Equal(t, "x", result[0].Y)
+}