From 87212466ae0c0f9ef90ea9d677c15ff413f40961 Mon Sep 17 00:00:00 2001
From: Sam Fredrickson <samfredrickson@gmail.com>
Date: Thu, 22 Aug 2024 00:53:48 -0700
Subject: [PATCH] Add iterators for a couple queues.

---
 go.mod            |  2 +-
 pq/lib.go         | 16 ++++++++++++++++
 pq/lib_test.go    | 24 ++++++++++++++++++++++++
 queue/lib.go      | 31 +++++++++++++++++++++++++++++++
 queue/lib_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 117 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 83f382d..a43544d 100644
--- a/go.mod
+++ b/go.mod
@@ -1,5 +1,5 @@
 module gogs.humancabbage.net/sam/priorityq
 
-go 1.20
+go 1.23
 
 require golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2
diff --git a/pq/lib.go b/pq/lib.go
index 8f45a8f..461c234 100644
--- a/pq/lib.go
+++ b/pq/lib.go
@@ -24,6 +24,7 @@
 package pq
 
 import (
+	"iter"
 	"sync"
 
 	"gogs.humancabbage.net/sam/priorityq"
@@ -153,3 +154,18 @@ func (s *state[P, T]) TrySend(priority P, value T) error {
 	s.canRecv.Broadcast()
 	return nil
 }
+
+// Iter returns an iterator that consumes values until the queue is closed.
+func (s *state[P, T]) Iter() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for {
+			_, t, ok := s.Recv()
+			if !ok {
+				return
+			}
+			if !yield(t) {
+				return
+			}
+		}
+	}
+}
diff --git a/pq/lib_test.go b/pq/lib_test.go
index 6dc1dc9..d64b209 100644
--- a/pq/lib_test.go
+++ b/pq/lib_test.go
@@ -147,6 +147,30 @@ func TestConcProducerConsumer(t *testing.T) {
 	wg.Wait()
 }
 
+func TestIter(t *testing.T) {
+	t.Parallel()
+	q := pq.Make[int, int](4)
+	q.Send(4, 0)
+	q.Send(3, 1)
+	q.Send(2, 2)
+	q.Send(1, 3)
+	q.Close()
+	i := 0
+	for v := range q.Iter() {
+		if v != i {
+			t.Errorf("expected %d, got %d", i, v)
+		}
+		i++
+	}
+
+	// to test yield() returning false
+	q = pq.Make[int, int](4)
+	q.Send(1, 3)
+	for _ = range q.Iter() {
+		break
+	}
+}
+
 func BenchmarkSend(b *testing.B) {
 	q := pq.Make[int, int](b.N)
 	// randomize priorities to get amortized cost per op
diff --git a/queue/lib.go b/queue/lib.go
index f64c2c4..5be5cec 100644
--- a/queue/lib.go
+++ b/queue/lib.go
@@ -6,6 +6,8 @@
 // makes determining whether the queue is full or empty trivial.
 package queue
 
+import "iter"
+
 // Q is a non-concurrent queue.
 type Q[T any] struct {
 	buf  []T
@@ -73,3 +75,32 @@ func (b *Q[T]) PushBack(value T) {
 		b.tail = 0
 	}
 }
+
+// Iter returns an iterator over all items in the queue.
+//
+// This does not pop items off the queue.
+func (b *Q[T]) Iter() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		remaining := b.len
+		i := b.head
+		for remaining > 0 {
+			yield(b.buf[i])
+			remaining--
+			i++
+			if i == b.len {
+				i = 0
+			}
+		}
+	}
+}
+
+// IterPop returns an iterator that pops each item off the queue.
+func (b *Q[T]) IterPop() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for b.CanPop() {
+			if !yield(b.PopFront()) {
+				break
+			}
+		}
+	}
+}
diff --git a/queue/lib_test.go b/queue/lib_test.go
index 90219af..c7e1c93 100644
--- a/queue/lib_test.go
+++ b/queue/lib_test.go
@@ -73,6 +73,51 @@ func TestFullPushPanic(t *testing.T) {
 	q.PushBack(2)
 }
 
+func TestIter(t *testing.T) {
+	q := queue.Make[int](4)
+	q.PushBack(1)
+	q.PushBack(2)
+	q.PushBack(3)
+	q.PushBack(4)
+	i := 0
+	for v := range q.Iter() {
+		expected := i + 1
+		if v != expected {
+			t.Errorf("iter %d should have value %d, not %d",
+				i, expected, v)
+		}
+		i++
+	}
+	if q.Len() != 4 {
+		t.Errorf("wrong length")
+	}
+}
+
+func TestIterPop(t *testing.T) {
+	t.Parallel()
+	q := queue.Make[int](4)
+	q.PushBack(1)
+	q.PushBack(2)
+	q.PushBack(3)
+	q.PushBack(4)
+	i := 0
+	for v := range q.IterPop() {
+		expected := i + 1
+		if v != expected {
+			t.Errorf("iter %d should have value %d, not %d",
+				i, expected, v)
+		}
+		i++
+		// to test yield() returning false
+		if i == 4 {
+			break
+		}
+	}
+	if q.Len() != 0 {
+		t.Errorf("wrong length")
+	}
+}
+
 func BenchmarkPush(b *testing.B) {
 	q := queue.Make[int](b.N)
 	rs := rand.NewSource(0)