move PQ to thirdparty

container.go

+package pq
+
+import "container/heap"
+
+// PQ is a basic priority queue.
+type PQ interface {
+	// Push adds the ele
+	Push(Elem)
+	// Pop returns the highest priority Elem in PQ.
+	Pop() Elem
+	// Len returns the number of elements in the PQ.
+	Len() int
+	// Update `fixes` the PQ.
+	Update(index int)
+
+	// TODO  explain why this interface should not be extended
+	// It does not support Remove. This is because...
+}
+
+// Elem describes elements that can be added to the PQ. Clients must implement
+// this interface.
+type Elem interface {
+	// SetIndex stores the int index.
+	SetIndex(int)
+	// Index returns the last given by SetIndex(int).
+	Index() int
+}
+
+// ElemComparator returns true if pri(a) > pri(b)
+type ElemComparator func(a, b Elem) bool
+
+// New creates a PQ with a client-supplied comparator.
+func New(cmp ElemComparator) PQ {
+	q := &wrapper{heapinterface{
+		elems: make([]Elem, 0),
+		cmp:   cmp,
+	}}
+	heap.Init(&q.heapinterface)
+	return q
+}
+
+// wrapper exists because we cannot re-define Push. We want to expose
+// Push(Elem) but heap.Interface requires Push(interface{})
+type wrapper struct {
+	heapinterface
+}
+
+var _ PQ = &wrapper{}
+
+func (w *wrapper) Push(e Elem) {
+	heap.Push(&w.heapinterface, e)
+}
+
+func (w *wrapper) Pop() Elem {
+	return heap.Pop(&w.heapinterface).(Elem)
+}
+
+func (w *wrapper) Update(index int) {
+	heap.Fix(&w.heapinterface, index)
+}
+
+// heapinterface handles dirty low-level details of managing the priority queue.
+type heapinterface struct {
+	elems []Elem
+	cmp   ElemComparator
+}
+
+var _ heap.Interface = &heapinterface{}
+
+// public interface
+
+func (q *heapinterface) Len() int {
+	return len(q.elems)
+}
+
+// Less delegates the decision to the comparator
+func (q *heapinterface) Less(i, j int) bool {
+	return q.cmp(q.elems[i], q.elems[j])
+}
+
+// Swap swaps the elements with indexes i and j.
+func (q *heapinterface) Swap(i, j int) {
+	q.elems[i], q.elems[j] = q.elems[j], q.elems[i]
+	q.elems[i].SetIndex(i)
+	q.elems[j].SetIndex(j)
+}
+
+// Note that Push and Pop in this interface are for package heap's
+// implementation to call. To add and remove things from the heap, wrap with
+// the pq struct to call heap.Push and heap.Pop.
+
+func (q *heapinterface) Push(x interface{}) { // where to put the elem?
+	t := x.(Elem)
+	t.SetIndex(len(q.elems))
+	q.elems = append(q.elems, t)
+}
+
+func (q *heapinterface) Pop() interface{} {
+	old := q.elems
+	n := len(old)
+	elem := old[n-1]       // remove the last
+	elem.SetIndex(-1)      // for safety // FIXME why?
+	q.elems = old[0 : n-1] // shrink
+	return elem
+}

container_test.go

+package pq
+
+import (
+	"sort"
+	"testing"
+)
+
+type TestElem struct {
+	Key      string
+	Priority int
+	index    int
+}
+
+func (e *TestElem) Index() int {
+	return e.index
+}
+
+func (e *TestElem) SetIndex(i int) {
+	e.index = i
+}
+
+var PriorityComparator = func(i, j Elem) bool {
+	return i.(*TestElem).Priority > j.(*TestElem).Priority
+}
+
+func TestQueuesReturnTypeIsSameAsParameterToPush(t *testing.T) {
+	q := New(PriorityComparator)
+	expectedKey := "foo"
+	elem := &TestElem{Key: expectedKey}
+	q.Push(elem)
+	switch v := q.Pop().(type) {
+	case *TestElem:
+		if v.Key != expectedKey {
+			t.Fatal("the key doesn't match the pushed value")
+		}
+	default:
+		t.Fatal("the queue is not casting values appropriately")
+	}
+}
+
+func TestCorrectnessOfPop(t *testing.T) {
+	q := New(PriorityComparator)
+	tasks := []TestElem{
+		TestElem{Key: "a", Priority: 9},
+		TestElem{Key: "b", Priority: 4},
+		TestElem{Key: "c", Priority: 3},
+		TestElem{Key: "d", Priority: 0},
+		TestElem{Key: "e", Priority: 6},
+	}
+	for _, e := range tasks {
+		q.Push(&e)
+	}
+	var priorities []int
+	for q.Len() > 0 {
+		i := q.Pop().(*TestElem).Priority
+		t.Log("popped %v", i)
+		priorities = append(priorities, i)
+	}
+	if !sort.IntsAreSorted(priorities) {
+		t.Fatal("the values were not returned in sorted order")
+	}
+}
+
+func TestUpdate(t *testing.T) {
+	t.Log(`
+	Add 3 elements.
+	Update the highest priority element to have the lowest priority and fix the queue.
+	It should come out last.`)
+	q := New(PriorityComparator)
+	lowest := &TestElem{Key: "originallyLowest", Priority: 1}
+	middle := &TestElem{Key: "originallyMiddle", Priority: 2}
+	highest := &TestElem{Key: "toBeUpdated", Priority: 3}
+	q.Push(middle)
+	q.Push(highest)
+	q.Push(lowest)
+	if q.Pop().(*TestElem).Key != highest.Key {
+		t.Fatal("popped element doesn't have the highest priority")
+	}
+	q.Push(highest)           // re-add the popped element
+	highest.Priority = 0      // update the PQ
+	q.Update(highest.Index()) // fix the PQ
+	if q.Pop().(*TestElem).Key != middle.Key {
+		t.Fatal("middle element should now have the highest priority")
+	}
+}