Add iterators

deque.go

@@ -1,6 +1,9 @@
 package deque
 
-import "fmt"
+import (
+	"fmt"
+	"iter"
+)
 
 // minCapacity is the smallest capacity that deque may have. Must be power of 2
 // for bitwise modulus: x % n == x & (n - 1).
@@ -94,6 +97,29 @@ func (q *Deque[T]) PopFront() T {
 	return ret
 }
 
+// IterPopFront returns an iterator the iteratively removes items from the
+// Front of the deque. This is more efficient than removing items one at a time
+// because it avoids intermediate resizing. If a resize is necessary, only one
+// is done when iteration ends.
+func (q *Deque[T]) IterPopFront() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		if q.Len() == 0 {
+			return
+		}
+		var zero T
+		for q.count != 0 {
+			ret := q.buf[q.head]
+			q.buf[q.head] = zero
+			q.head = q.next(q.head)
+			q.count--
+			if !yield(ret) {
+				break
+			}
+		}
+		q.shrinkToFit()
+	}
+}
+
 // PopBack removes and returns the element from the back of the queue.
 // Implements LIFO when used with PushBack. If the queue is empty, the call
 // panics.
@@ -115,6 +141,29 @@ func (q *Deque[T]) PopBack() T {
 	return ret
 }
 
+// IterPopBack returns an iterator the iteratively removes items from the back
+// of the deque. This is more efficient than removing items one at a time
+// because it avoids intermediate resizing. If a resize is necessary, only one
+// is done when iteration ends.
+func (q *Deque[T]) IterPopBack() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		if q.Len() == 0 {
+			return
+		}
+		var zero T
+		for q.count != 0 {
+			q.tail = q.prev(q.tail)
+			ret := q.buf[q.tail]
+			q.buf[q.tail] = zero
+			q.count--
+			if !yield(ret) {
+				break
+			}
+		}
+		q.shrinkToFit()
+	}
+}
+
 // Front returns the element at the front of the queue. This is the element
 // that would be returned by PopFront. This call panics if the queue is empty.
 func (q *Deque[T]) Front() T {
@@ -160,6 +209,50 @@ func (q *Deque[T]) Set(i int, item T) {
 	q.buf[(q.head+i)&(len(q.buf)-1)] = item
 }
 
+// Iter returns a go iterator to range over all items in the Deque, yielding
+// the index of each item and the item, from front to back. Modification of
+// Deque during iteration panics.
+func (q *Deque[T]) Iter() iter.Seq2[int, T] {
+	return func(yield func(int, T) bool) {
+		if q.Len() == 0 {
+			return
+		}
+		count := q.count
+		head := q.head
+		for i := 0; i < count; i++ {
+			if q.count != count {
+				panic("deque: modified during iteration")
+			}
+			if !yield(i, q.buf[head]) {
+				return
+			}
+			head = q.next(head)
+		}
+	}
+}
+
+// RIter returns a go iterator to range over all items in the Deque, yielding
+// the index of each item and the item, from back to front. Modification of
+// Deque during iteration panics.
+func (q *Deque[T]) RIter() iter.Seq2[int, T] {
+	return func(yield func(int, T) bool) {
+		if q.Len() == 0 {
+			return
+		}
+		count := q.count
+		tail := q.tail
+		for i := count - 1; i >= 0; i-- {
+			if q.count != count {
+				panic("deque: modified during iteration")
+			}
+			tail = q.prev(tail)
+			if !yield(i, q.buf[tail]) {
+				return
+			}
+		}
+	}
+}
+
 // Clear removes all elements from the queue, but retains the current capacity.
 // This is useful when repeatedly reusing the queue at high frequency to avoid
 // GC during reuse. The queue will not be resized smaller as long as items are
@@ -416,6 +509,23 @@ func (q *Deque[T]) shrinkIfExcess() {
 	}
 }
 
+func (q *Deque[T]) shrinkToFit() {
+	if len(q.buf) > q.minCap && (q.count<<2) <= len(q.buf) {
+		if q.count == 0 {
+			q.head = 0
+			q.tail = 0
+			q.buf = make([]T, minCapacity)
+			return
+		}
+
+		c := minCapacity
+		for c < q.count {
+			c <<= 1
+		}
+		q.resize(c)
+	}
+}
+
 // resize resizes the deque to fit exactly twice its current contents. This is
 // used to grow the queue when it is full, and also to shrink it when it is
 // only a quarter full.

deque_test.go

@@ -384,14 +384,14 @@ func TestAt(t *testing.T) {
 	// Front to back.
 	for j := 0; j < q.Len(); j++ {
 		if q.At(j) != j {
-			t.Errorf("index %d doesn't contain %d", j, j)
+			t.Errorf("wrong item at index %d", j)
 		}
 	}
 
 	// Back to front
-	for j := 1; j <= q.Len(); j++ {
-		if q.At(q.Len()-j) != q.Len()-j {
-			t.Errorf("index %d doesn't contain %d", q.Len()-j, q.Len()-j)
+	for j := q.Len() - 1; j >= 0; j-- {
+		if q.At(j) != j {
+			t.Errorf("wrong item at index %d", j)
 		}
 	}
 }
@@ -674,6 +674,227 @@ func TestSwap(t *testing.T) {
 	})
 }
 
+func TestIter(t *testing.T) {
+	var q Deque[int]
+
+	for range q.Iter() {
+		t.Fatal("iterated when empty")
+	}
+
+	q.Grow(50)
+	for i := 0; i < 50; i++ {
+		q.PushBack(i)
+	}
+
+	// Front to back.
+	expect := 0
+	for i, item := range q.Iter() {
+		if i != expect {
+			t.Fatalf("expected index %d, got %d", expect, i)
+		}
+		if item != i {
+			t.Errorf("index %d contains %d", i, item)
+		}
+		if i == 40 {
+			break
+		}
+		expect++
+	}
+
+	assertPanics(t, "Iter must panic when deque modified during iteration", func() {
+		for i, _ := range q.Iter() {
+			if i == 42 {
+				q.PushBack(51)
+			}
+		}
+	})
+}
+
+func TestRIter(t *testing.T) {
+	var q Deque[int]
+
+	for range q.RIter() {
+		t.Fatal("iterated when empty")
+	}
+
+	q.Grow(50)
+	for i := 0; i < 50; i++ {
+		q.PushBack(i)
+	}
+
+	// Back to fron
+	expect := 49
+	for i, item := range q.RIter() {
+		if i != expect {
+			t.Fatalf("expected index %d, got %d", expect, i)
+		}
+		if item != i {
+			t.Fatalf("index %d contains %d", i, item)
+		}
+		if i == 10 {
+			break
+		}
+		expect--
+	}
+
+	assertPanics(t, "RIter must panic when deque modified during iteration", func() {
+		for i, _ := range q.RIter() {
+			if i == 42 {
+				q.PushBack(51)
+			}
+		}
+	})
+}
+
+func TestIterPopBack(t *testing.T) {
+	var q Deque[int]
+	size := minCapacity * 5
+
+	q.Grow(size)
+	for i := 0; i < size; i++ {
+		q.PushBack(i)
+	}
+
+	last := q.Front()
+	var removed, lastRm int
+	for i := range q.IterPopBack() {
+		removed++
+		lastRm = i
+	}
+
+	if last != lastRm {
+		t.Fatal("did not expose expected item list")
+	}
+	if removed != size {
+		t.Fatal("wrong removed count")
+	}
+	if lastRm != 0 {
+		t.Fatal("wrong last item removed")
+	}
+	if q.Len() != 0 {
+		t.Error("q.Len() =", q.Len(), "expected 0")
+	}
+
+	q.Grow(size)
+	for i := 0; i < size; i++ {
+		q.PushBack(i)
+	}
+	last = q.Front()
+	removed = 0
+	for i := range q.IterPopBack() {
+		removed++
+		lastRm = i
+	}
+	if last != lastRm {
+		t.Fatal("did not expose expected item list")
+	}
+	if removed != size {
+		t.Fatal("wrong removed count, got", removed, " expected", size)
+	}
+	if q.Len() != 0 {
+		t.Error("q.Len() =", q.Len(), "expected 0")
+	}
+	for range q.IterPopBack() {
+		t.Fatal("iteration with 0 items")
+	}
+
+	for i := 0; i < 5; i++ {
+		q.PushBack(i)
+	}
+	c := q.Cap()
+	removed = 0
+	for range q.IterPopBack() {
+		removed++
+	}
+	if removed != 5 {
+		t.Fatal("wrong removed count")
+	}
+	if q.Cap() != c {
+		t.Fatal("unexpected capacity change")
+	}
+
+	for i := 0; i < 65; i++ {
+		q.PushBack(i)
+	}
+	c = q.Cap()
+	removed = 0
+	for range q.IterPopBack() {
+		removed++
+		if removed == 3 {
+			break
+		}
+	}
+	if q.Cap() != c {
+		t.Fatal("unexpected capacity change")
+	}
+}
+
+func TestIterPopFront(t *testing.T) {
+	var q Deque[int]
+	size := minCapacity * 5
+
+	q.Grow(size)
+	for i := 0; i < size; i++ {
+		q.PushBack(i)
+	}
+	last := q.Back()
+	var lastRm, removed int
+	for i := range q.IterPopFront() {
+		removed++
+		lastRm = i
+	}
+	if last != lastRm {
+		t.Fatal("did not expose expected item list")
+	}
+	if removed != size {
+		t.Fatal("wrong removed count")
+	}
+	if lastRm != size-1 {
+		t.Fatal("wrong last item removed")
+	}
+	if q.Len() != 0 {
+		t.Error("q.Len() =", q.Len(), "expected 0")
+	}
+	for range q.IterPopFront() {
+		t.Fatal("iteration with 0 items")
+	}
+
+	for i := 0; i < 5; i++ {
+		q.PushBack(i)
+	}
+	c := q.Cap()
+	removed = 0
+	for range q.IterPopFront() {
+		removed++
+	}
+	if removed != 5 {
+		t.Fatal("wrong removed count")
+	}
+	if q.Cap() != c {
+		t.Fatal("unexpected capacity change")
+	}
+
+	for i := 0; i < 65; i++ {
+		q.PushBack(i)
+	}
+	c = q.Cap()
+	removed = 0
+	for range q.IterPopFront() {
+		removed++
+		if removed == 3 {
+			break
+		}
+	}
+	for range q.IterPopFront() {
+		if q.Len() == 32 {
+			break
+		}
+	}
+	if q.Cap() == c {
+		t.Fatal("expected capacity change")
+	}
+}
+
 func TestFrontBackOutOfRangePanics(t *testing.T) {
 	const msg = "should panic when peeking empty queue"
 	var q Deque[int]

go.mod

@@ -1,3 +1,3 @@
 module github.com/gammazero/deque
 
-go 1.22
+go 1.23