From c1074bfa47a450b7ae67251f74d912fc109626bb Mon Sep 17 00:00:00 2001
From: Alexander <bombblob@gmail.com>
Date: Fri, 15 Sep 2023 23:22:56 -0700
Subject: [PATCH] Add channel support

---
 README.md        | 12 ++++++++++++
 iterator.go      | 39 ++++++++++++++++++++++++++++++++++++++-
 iterator_test.go | 23 +++++++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d0ca963..1559379 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,12 @@ func Slice[T any](slice []T) Iterator[T]
 
 `Slice` returns an Iterator that yields elements from a slice.
 
+```go
+func Chan[T any](ch <-chan T) Iterator[T]
+```
+
+`Chan` returns an Iterator that yields elements from a channel.
+
 ```go
 func String(input string) Iterator[rune]
 ```
@@ -151,6 +157,12 @@ func ToSlice[T any](it Iterator[T]) []T
 
 `ToSlice` consumes an Iterator creating a slice from the yielded values.
 
+```go
+func ToChan[T any](it Iterator[T]) <-chan T
+```
+
+`ToChan` consumes an Iterator writing the yielded values to a channel.
+
 ```go
 func ToString(it Iterator[rune]) string
 ```
diff --git a/iterator.go b/iterator.go
index b0198d5..80dd5f0 100644
--- a/iterator.go
+++ b/iterator.go
@@ -1,6 +1,8 @@
 package iter
 
-import "unicode/utf8"
+import (
+	"unicode/utf8"
+)
 
 // Iterator[T] represents an iterator yielding elements of type T.
 type Iterator[T any] interface {
@@ -77,6 +79,25 @@ func (it *sliceIter[T]) Next() Option[T] {
 	return Some[T](first)
 }
 
+type chanIter[T any] struct {
+	ch <-chan T
+}
+
+// Chan returns an Iterator that yields elements written to a channel.
+func Chan[T any](ch <-chan T) Iterator[T] {
+	return &chanIter[T]{
+		ch: ch,
+	}
+}
+
+func (c *chanIter[T]) Next() Option[T] {
+	v, ok := <-c.ch
+	if !ok {
+		return None[T]()
+	}
+	return Some(v)
+}
+
 // ToSlice consumes an Iterator creating a slice from the yielded values.
 func ToSlice[T any](it Iterator[T]) []T {
 	result := []T{}
@@ -91,6 +112,22 @@ func ToString(it Iterator[rune]) string {
 	return string(ToSlice(it))
 }
 
+// ToChan consumes an Iterator writing yielded values to the returned channel.
+func ToChan[T any](it Iterator[T]) <-chan T {
+	ch := make(chan T)
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+			}
+		}()
+		ForEach(it, func(v T) {
+			ch <- v
+		})
+		close(ch)
+	}()
+	return ch
+}
+
 type mapIter[T, R any] struct {
 	inner Iterator[T]
 	fn    func(T) R
diff --git a/iterator_test.go b/iterator_test.go
index 85feeb3..d6fd660 100644
--- a/iterator_test.go
+++ b/iterator_test.go
@@ -19,6 +19,22 @@ func TestSlice(t *testing.T) {
 	equals(t, it.Next().IsNone(), true)
 }
 
+func TestChan(t *testing.T) {
+	ch := make(chan int)
+	arr := []int{1, 2, 3}
+	go func() {
+		for _, v := range arr {
+			ch <- v
+		}
+		close(ch)
+	}()
+	it := Chan(ch)
+	equals(t, it.Next().Unwrap(), 1)
+	equals(t, it.Next().Unwrap(), 2)
+	equals(t, it.Next().Unwrap(), 3)
+	equals(t, it.Next().IsNone(), true)
+}
+
 func TestRepeat(t *testing.T) {
 	it := Repeat(5)
 	equals(t, it.Next().Unwrap(), 5)
@@ -75,6 +91,13 @@ func TestToSlice(t *testing.T) {
 	equals(t, slice2, []int{})
 }
 
+func TestToChan(t *testing.T) {
+	slice1 := ToSlice(Chan(ToChan(Take(Repeat(5), 3))))
+	equals(t, slice1, []int{5, 5, 5})
+	slice2 := ToSlice(Chan(ToChan(Empty[int]())))
+	equals(t, slice2, []int{})
+}
+
 func TestDrop(t *testing.T) {
 	slice1 := ToSlice(Drop(Slice([]int{1, 2, 3, 4, 5}), 3))
 	equals(t, slice1, []int{4, 5})