feature/add generic types

go.mod

@@ -0,0 +1,3 @@
+module github.com/airspace-link-inc/go-priority-queue
+
+go 1.18

priorty_queue.go

@@ -11,33 +11,34 @@ import (
 	"errors"
 )
 
-// PriorityQueue represents the queue
-type PriorityQueue struct {
-	itemHeap *itemHeap
-	lookup   map[interface{}]*item
+var ErrEmptyQueue = errors.New("empty queue")
+
+// PriorityQueue[T] represents the queue of comparable type T.
+type PriorityQueue[T comparable] struct {
+	itemHeap *itemHeap[T]
+	lookup   map[T]*item[T]
 }
 
-// New initializes an empty priority queue.
-func New() PriorityQueue {
-	return PriorityQueue{
-		itemHeap: &itemHeap{},
-		lookup:   make(map[interface{}]*item),
+// New[T] initializes an empty priority queue of type T.
+func New[T comparable]() PriorityQueue[T] {
+	return PriorityQueue[T]{
+		itemHeap: &itemHeap[T]{},
+		lookup:   make(map[T]*item[T]),
 	}
 }
 
 // Len returns the number of elements in the queue.
-func (p *PriorityQueue) Len() int {
+func (p *PriorityQueue[T]) Len() int {
 	return p.itemHeap.Len()
 }
 
 // Insert inserts a new element into the queue. No action is performed on duplicate elements.
-func (p *PriorityQueue) Insert(v interface{}, priority float64) {
-	_, ok := p.lookup[v]
-	if ok {
+func (p *PriorityQueue[T]) Insert(v T, priority float64) {
+	if _, ok := p.lookup[v]; ok {
 		return
 	}
 
-	newItem := &item{
+	newItem := &item[T]{
 		value:    v,
 		priority: priority,
 	}
@@ -47,19 +48,20 @@ func (p *PriorityQueue) Insert(v interface{}, priority float64) {
 
 // Pop removes the element with the highest priority from the queue and returns it.
 // In case of an empty queue, an error is returned.
-func (p *PriorityQueue) Pop() (interface{}, error) {
+func (p *PriorityQueue[T]) Pop() (T, error) {
 	if len(*p.itemHeap) == 0 {
-		return nil, errors.New("empty queue")
+		var zeroVal T
+		return zeroVal, ErrEmptyQueue
 	}
 
-	item := heap.Pop(p.itemHeap).(*item)
+	item := heap.Pop(p.itemHeap).(*item[T])
 	delete(p.lookup, item.value)
 	return item.value, nil
 }
 
 // UpdatePriority changes the priority of a given item.
 // If the specified item is not present in the queue, no action is performed.
-func (p *PriorityQueue) UpdatePriority(x interface{}, newPriority float64) {
+func (p *PriorityQueue[T]) UpdatePriority(x T, newPriority float64) {
 	item, ok := p.lookup[x]
 	if !ok {
 		return
@@ -69,35 +71,35 @@ func (p *PriorityQueue) UpdatePriority(x interface{}, newPriority float64) {
 	heap.Fix(p.itemHeap, item.index)
 }
 
-type itemHeap []*item
+type itemHeap[T comparable] []*item[T]
 
-type item struct {
-	value    interface{}
+type item[T comparable] struct {
+	value    T
 	priority float64
 	index    int
 }
 
-func (ih *itemHeap) Len() int {
+func (ih *itemHeap[T]) Len() int {
 	return len(*ih)
 }
 
-func (ih *itemHeap) Less(i, j int) bool {
+func (ih *itemHeap[T]) Less(i, j int) bool {
 	return (*ih)[i].priority < (*ih)[j].priority
 }
 
-func (ih *itemHeap) Swap(i, j int) {
+func (ih *itemHeap[T]) Swap(i, j int) {
 	(*ih)[i], (*ih)[j] = (*ih)[j], (*ih)[i]
 	(*ih)[i].index = i
 	(*ih)[j].index = j
 }
 
-func (ih *itemHeap) Push(x interface{}) {
-	it := x.(*item)
+func (ih *itemHeap[T]) Push(x interface{}) {
+	it := x.(*item[T])
 	it.index = len(*ih)
 	*ih = append(*ih, it)
 }
 
-func (ih *itemHeap) Pop() interface{} {
+func (ih *itemHeap[T]) Pop() interface{} {
 	old := *ih
 	item := old[len(old)-1]
 	*ih = old[0 : len(old)-1]

priorty_queue_test.go

@@ -1,12 +1,13 @@
 package pq
 
 import (
+	"math"
 	"sort"
 	"testing"
 )
 
 func TestPriorityQueue(t *testing.T) {
-	pq := New()
+	pq := New[float64]()
 	elements := []float64{5, 3, 7, 8, 6, 2, 9}
 	for _, e := range elements {
 		pq.Insert(e, e)
@@ -19,15 +20,15 @@ func TestPriorityQueue(t *testing.T) {
 			t.Fatalf(err.Error())
 		}
 
-		i := item.(float64)
+		i := item
 		if e != i {
 			t.Fatalf("expected %v, got %v", e, i)
 		}
 	}
 }
 
 func TestPriorityQueueUpdate(t *testing.T) {
-	pq := New()
+	pq := New[string]()
 	pq.Insert("foo", 3)
 	pq.Insert("bar", 4)
 	pq.UpdatePriority("bar", 2)
@@ -37,13 +38,13 @@ func TestPriorityQueueUpdate(t *testing.T) {
 		t.Fatal(err.Error())
 	}
 
-	if item.(string) != "bar" {
+	if item != "bar" {
 		t.Fatal("priority update failed")
 	}
 }
 
 func TestPriorityQueueLen(t *testing.T) {
-	pq := New()
+	pq := New[string]()
 	if pq.Len() != 0 {
 		t.Fatal("empty queue should have length of 0")
 	}
@@ -55,8 +56,24 @@ func TestPriorityQueueLen(t *testing.T) {
 	}
 }
 
+func TestItemHeapLess(t *testing.T) {
+	h := itemHeap[int]{
+		&item[int]{priority: math.Inf(1)},
+		&item[int]{priority: math.Inf(1)},
+		&item[int]{priority: math.Inf(-1)},
+		&item[int]{priority: math.Inf(-1)},
+	}
+
+	// test all pairwise elements: 0,1 then 2,3 then...
+	for i := 0; i < len(h); i += 2 {
+		if h.Less(i, i+1) {
+			t.Fatalf("%v should not have less priority than %v", h[i], h[i+1])
+		}
+	}
+}
+
 func TestDoubleAddition(t *testing.T) {
-	pq := New()
+	pq := New[string]()
 	pq.Insert("foo", 2)
 	pq.Insert("bar", 3)
 	pq.Insert("bar", 1)
@@ -66,21 +83,23 @@ func TestDoubleAddition(t *testing.T) {
 	}
 
 	item, _ := pq.Pop()
-	if item.(string) != "foo" {
+	if item != "foo" {
 		t.Fatal("queue should ignore duplicate insert, not update existing item")
 	}
 }
 
 func TestPopEmptyQueue(t *testing.T) {
-	pq := New()
+	pq := New[float32]()
 	_, err := pq.Pop()
 	if err == nil {
 		t.Fatal("should produce error when performing pop on empty queue")
+	} else if err != ErrEmptyQueue {
+		t.Fatalf("error should be equal to %v, got %v", ErrEmptyQueue, err)
 	}
 }
 
 func TestUpdateNonExistingItem(t *testing.T) {
-	pq := New()
+	pq := New[string]()
 
 	pq.Insert("foo", 4)
 	pq.UpdatePriority("bar", 5)
@@ -90,7 +109,7 @@ func TestUpdateNonExistingItem(t *testing.T) {
 	}
 
 	item, _ := pq.Pop()
-	if item.(string) != "foo" {
-		t.Fatalf("update should not overwrite item, expected \"foo\", got \"%v\"", item.(string))
+	if item != "foo" {
+		t.Fatalf("update should not overwrite item, expected \"foo\", got \"%v\"", item)
 	}
 }