Just deque

1 files changed, 211 insertions, 0 deletions

diff --git a/deque.go b/deque.go
new file mode 100644
index 0000000..e8a9018
--- /dev/null
+++ b/deque.go
@@ -0,0 +1,211 @@
+// Package deque implements a double-ended queue (deque) implemented with a
+// slice-backed ring buffer.
+//
+// This queue has O(1) amortized inserts and removals from both ends of the
+// container. It also has O(1) indexing like a vector.
+//
+// The "default" usage of this type as a queue is to use [Deque.PushBack] to add
+// to the queue, and [Deque.PopFront] to remove from the queue. Iterating over
+// Deque goes front to back.
+//
+// The core implementation is "ported" (stolen) from Rust's VecDeque.
+package deque
+
+import (
+	"fmt"
+	"iter"
+	"slices"
+	"strings"
+)
+
+// Deque is a double-ended queue. The zero value is ready for use.
+type Deque[T any] struct {
+	head int
+	buf  []T
+}
+
+// WithCapacity allocates a deque with the given capacity.
+func WithCapacity[T any](cap int) *Deque[T] {
+	return &Deque[T]{buf: make([]T, 0, cap)}
+}
+
+// From creates a new queue using the given slice as the backing buffer.
+func From[S ~[]T, T any](slice S) *Deque[T] {
+	return &Deque[T]{buf: slice}
+}
+
+func (q *Deque[T]) wrapAdd(i, addend int) int {
+	i += addend
+	if i >= cap(q.buf) {
+		return i - cap(q.buf)
+	}
+	return i
+}
+
+func (q *Deque[T]) toPhysicalIdx(i int) int {
+	return q.wrapAdd(q.head, i)
+}
+
+// At returns the item at position i. At panics if i < 0 or i >= q.Len().
+func (q *Deque[T]) At(i int) T {
+	if !(0 <= i && i < len(q.buf)) {
+		panic(fmt.Sprintf("index out of range [%d] with length %d", i, len(q.buf)))
+	}
+	return q.buf[:cap(q.buf)][q.toPhysicalIdx(i)]
+}
+
+// Cap returns the number of elements the deque can hold without reallocating.
+func (q *Deque[T]) Cap() int {
+	return cap(q.buf)
+}
+
+// Len returns the number of elements in the deque.
+func (q *Deque[T]) Len() int {
+	return len(q.buf)
+}
+
+// PopFront removes and returns the item at index 0 if the deque is non-empty.
+func (q *Deque[T]) PopFront() (T, bool) {
+	if len(q.buf) == 0 {
+		var zero T
+		return zero, false
+	}
+	oldHead := q.head
+	q.head = q.toPhysicalIdx(1)
+	q.buf = q.buf[:len(q.buf)-1]
+	return q.buf[:cap(q.buf)][oldHead], true
+}
+
+// PopBack removes and returns the last item in the deque if it is non-empty.
+func (q *Deque[T]) PopBack() (T, bool) {
+	if len(q.buf) == 0 {
+		var zero T
+		return zero, false
+	}
+	q.buf = q.buf[:len(q.buf)-1]
+	return q.buf[:cap(q.buf)][q.toPhysicalIdx(len(q.buf))], true
+}
+
+// PushFront prepends the given items to the front of the deque.
+func (q *Deque[T]) PushFront(values ...T) {
+	q.Grow(len(values))
+	q.buf = q.buf[:len(q.buf)+len(values)]
+	if q.head >= len(values) {
+		newHead := q.head - len(values)
+		copy(q.buf[newHead:q.head], values)
+		q.head = newHead
+	} else {
+		tailLen := len(values) - q.head
+		copy(q.buf[:q.head], values[tailLen:])
+		copy(q.buf[cap(q.buf)-tailLen:cap(q.buf)], values[:tailLen])
+		q.head = cap(q.buf) - tailLen
+	}
+}
+
+// PushBack appends the given items to the back of the deque.
+func (q *Deque[T]) PushBack(values ...T) {
+	q.Grow(len(values))
+	endIdx := q.wrapAdd(q.head, len(q.buf))
+	if len(values) <= cap(q.buf)-endIdx {
+		copy(q.buf[endIdx:endIdx+len(values)], values)
+	} else {
+		headLen := cap(q.buf) - endIdx
+		copy(q.buf[endIdx:cap(q.buf)], values[:headLen])
+		copy(q.buf[:len(values)-headLen], values[headLen:])
+	}
+	q.buf = q.buf[:len(q.buf)+len(values)]
+}
+
+// Reset empties the deque, retaining the underlying storage for use by
+// future pushes.
+func (q *Deque[T]) Reset() {
+	q.buf = q.buf[:0]
+}
+
+// Grow makes space for at least n more elements to be inserted in the given
+// deque without reallocation.
+func (q *Deque[T]) Grow(n int) {
+	if n <= cap(q.buf)-len(q.buf) {
+		return
+	}
+
+	oldCap := cap(q.buf)
+	q.buf = slices.Grow(q.buf, n)
+	newCap := cap(q.buf)
+
+	// Move the shortest contiguous section of the ring buffer
+	//
+	// H := head
+	// L := last element (`self.to_physical_idx(self.len - 1)`)
+	//
+	//    H             L
+	//   [o o o o o o o o ]
+	//    H             L
+	// A [o o o o o o o o . . . . . . . . ]
+	//        L H
+	//   [o o o o o o o o ]
+	//          H             L
+	// B [. . . o o o o o o o o . . . . . ]
+	//              L H
+	//   [o o o o o o o o ]
+	//              L                 H
+	// C [o o o o o o . . . . . . . . o o ]
+
+	if q.head <= oldCap-len(q.buf) {
+		// A
+		return
+	}
+	headLen := oldCap - q.head
+	tailLen := len(q.buf) - headLen
+	if headLen > tailLen && newCap-oldCap >= tailLen {
+		// B
+		copy(q.buf[oldCap:oldCap+tailLen], q.buf[:tailLen])
+		return
+	}
+	// C
+	newHead := newCap - headLen
+	copy(q.buf[newHead:newHead+headLen], q.buf[q.head:q.head+headLen])
+	q.head = newHead
+}
+
+// All returns an iterator over the elements in the deque. It does not pop
+// any elements.
+func (q *Deque[T]) All() iter.Seq2[int, T] {
+	return func(yield func(int, T) bool) {
+		// Don't use range over int in case the length changes while
+		// we're iterating
+		for i := 0; i < len(q.buf); i++ {
+			if !yield(i, q.buf[:cap(q.buf)][q.toPhysicalIdx(i)]) {
+				return
+			}
+		}
+	}
+}
+
+// PopAll empties the deque and returns an iterator over the popped elements.
+// It's not safe to modify the deque while iterating using PopAll.
+func (q *Deque[T]) PopAll() iter.Seq[T] {
+	n := len(q.buf)
+	q.buf = q.buf[:0]
+	return func(yield func(T) bool) {
+		for i := range n {
+			if !yield(q.buf[:cap(q.buf)][q.toPhysicalIdx(i)]) {
+				return
+			}
+		}
+	}
+}
+
+// String displays the deque as a string, using fmt.Sprint to show each element.
+func (q *Deque[T]) String() string {
+	buf := new(strings.Builder)
+	buf.WriteString("[")
+	for i := range len(q.buf) {
+		if i > 0 {
+			buf.WriteString(" ")
+		}
+		fmt.Fprint(buf, q.buf[:cap(q.buf)][q.toPhysicalIdx(i)])
+	}
+	buf.WriteString("]")
+	return buf.String()
+}