slicehelpers: Add Stack

go.mod

@@ -2,12 +2,14 @@ module github.com/bep/helpers
 
 go 1.18
 
-require github.com/frankban/quicktest v1.14.3
+require (
+	github.com/frankban/quicktest v1.14.3
+	golang.org/x/sync v0.7.0
+)
 
 require (
 	github.com/google/go-cmp v0.5.8 // indirect
 	github.com/kr/pretty v0.3.0 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/rogpeppe/go-internal v1.8.1 // indirect
-	golang.org/x/sync v0.7.0 // indirect
 )

slicehelpers/slicehelpers.go

@@ -3,6 +3,8 @@
 
 package slicehelpers
 
+import "sync"
+
 // Chunk splits the slice s into n number of chunks.
 func Chunk[T any](s []T, n int) [][]T {
 	if len(s) == 0 {
@@ -25,7 +27,7 @@
 	return partitions
 }
 
-// Parition partitions s into slices of size size.
+// Partition partitions s into slices of size size.
 func Partition[T any](s []T, size int) [][]T {
 	if len(s) == 0 {
 		return nil
@@ -43,3 +45,110 @@
 	}
 	return partitions
 }
+
+// stack represents a stack data structure.
+type stack[T any] struct {
+	zero  T
+	items []T
+}
+
+type StackConfig struct {
+	// ThreadSafe indicates if the stack should be thread safe.
+	ThreadSafe bool
+}
+
+// NewStack returns a new Stack.
+func NewStack[T any](conf StackConfig) Stack[T] {
+	s := &stack[T]{}
+	if !conf.ThreadSafe {
+		return s
+	}
+	return &threadSafeStack[T]{stack: s}
+}
+
+// Push adds an element to the stack.
+func (s *stack[T]) Push(v T) {
+	s.items = append(s.items, v)
+}
+
+// Peek returns the top element of the stack.
+func (s *stack[T]) Peek() T {
+	if len(s.items) == 0 {
+		return s.zero
+	}
+	return s.items[len(s.items)-1]
+}
+
+// Pop removes and returns the top element of the stack.
+func (s *stack[T]) Pop() T {
+	if len(s.items) == 0 {
+		return s.zero
+	}
+	v := s.items[len(s.items)-1]
+	s.items = s.items[:len(s.items)-1]
+	return v
+}
+
+// Drain removes all elements from the stack and returns them.
+func (s *stack[T]) Drain() []T {
+	items := s.items
+	s.items = nil
+	return items
+}
+
+// Len returns the number of elements in the stack.
+func (s *stack[T]) Len() int {
+	return len(s.items)
+}
+
+type Stack[T any] interface {
+	// Push adds an element to the stack.
+	Push(v T)
+	// Pop removes and returns the top element of the stack.
+	Pop() T
+	// Peek returns the top element of the stack.
+	Peek() T
+	// Drain removes all elements from the stack and returns them.
+	Drain() []T
+	// Len returns the number of elements in the stack.
+	Len() int
+}
+
+type threadSafeStack[T any] struct {
+	*stack[T]
+	mu sync.RWMutex
+}
+
+func (s *threadSafeStack[T]) Push(v T) {
+	s.mu.Lock()
+	s.stack.Push(v)
+	s.mu.Unlock()
+}
+
+func (s *threadSafeStack[T]) Pop() T {
+	s.mu.Lock()
+	v := s.stack.Pop()
+	s.mu.Unlock()
+	return v
+}
+
+func (s *threadSafeStack[T]) Peek() T {
+	s.mu.RLock()
+	v := s.stack.Peek()
+	s.mu.RUnlock()
+	return v
+}
+
+func (s *threadSafeStack[T]) Drain() []T {
+	s.mu.Lock()
+	items := s.stack.Drain()
+	s.mu.Unlock()
+	return items
+}
+
+func (s *threadSafeStack[T]) Len() int {
+	s.mu.RLock()
+	l := s.stack.Len()
+	s.mu.RUnlock()
+	return l
+}

slicehelpers/slicehelpers_test.go

@@ -4,6 +4,7 @@
 package slicehelpers
 
 import (
+	"sync"
 	"testing"
 
 	qt "github.com/frankban/quicktest"
@@ -135,3 +136,66 @@ func TestPartition(t *testing.T) {
 		qt.IsNil,
 	)
 }
+
+func TestStack(t *testing.T) {
+	c := qt.New(t)
+
+	s := NewStack[int](StackConfig{})
+
+	c.Assert(s.Peek(), qt.Equals, 0)
+	c.Assert(s.Pop(), qt.Equals, 0)
+	s.Push(1)
+	c.Assert(s.Peek(), qt.Equals, 1)
+	c.Assert(s.Pop(), qt.Equals, 1)
+	c.Assert(s.Pop(), qt.Equals, 0)
+	c.Assert(s.Peek(), qt.Equals, 0)
+
+	s.Push(2)
+	s.Push(3)
+	c.Assert(s.Len(), qt.Equals, 2)
+	c.Assert(s.Peek(), qt.Equals, 3)
+	c.Assert(s.Pop(), qt.Equals, 3)
+	c.Assert(s.Pop(), qt.Equals, 2)
+	c.Assert(s.Pop(), qt.Equals, 0)
+	c.Assert(s.Peek(), qt.Equals, 0)
+
+	s.Push(4)
+	s.Push(5)
+	c.Assert(s.Drain(), qt.DeepEquals, []int{4, 5})
+	c.Assert(s.Len(), qt.Equals, 0)
+}
+
+func TestStackThreadSafe(t *testing.T) {
+	s := NewStack[int](StackConfig{ThreadSafe: true})
+
+	var wg sync.WaitGroup
+
+	for k := 0; k < 20; k++ {
+		wg.Add(3)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 100; i++ {
+				s.Push(i)
+				s.Len()
+				s.Peek()
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 50; i++ {
+				s.Pop()
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 50; i++ {
+				s.Drain()
+			}
+		}()
+
+	}
+
+	wg.Wait()
+}