github.com/anycable/anycable-go@v1.5.1/utils/priority_queue.go (about)

     1  package utils
     2  
     3  // Based on https://pkg.go.dev/container/heap
     4  
     5  import (
     6  	"container/heap"
     7  
     8  	"golang.org/x/exp/constraints"
     9  )
    10  
    11  // An PriorityQueueItem is something we manage in a priority queue.
    12  type PriorityQueueItem[T any, P constraints.Ordered] struct {
    13  	value    T
    14  	priority P
    15  	// The index is needed by update and is maintained by the heap.Interface methods.
    16  	index int
    17  }
    18  
    19  func (pq PriorityQueueItem[T, P]) Value() T {
    20  	return pq.value
    21  }
    22  
    23  func (pq PriorityQueueItem[T, P]) Priority() P {
    24  	return pq.priority
    25  }
    26  
    27  type PriorityQueue[T any, P constraints.Ordered] []*PriorityQueueItem[T, P]
    28  
    29  func NewPriorityQueue[T any, P constraints.Ordered]() *PriorityQueue[T, P] {
    30  	pq := &PriorityQueue[T, P]{}
    31  	heap.Init(pq)
    32  	return pq
    33  }
    34  
    35  func (pq *PriorityQueue[T, P]) PushItem(v T, priority P) *PriorityQueueItem[T, P] {
    36  	item := &PriorityQueueItem[T, P]{value: v, priority: priority}
    37  	heap.Push(pq, item)
    38  	return item
    39  }
    40  
    41  func (pq *PriorityQueue[T, P]) PopItem() *PriorityQueueItem[T, P] {
    42  	return heap.Pop(pq).(*PriorityQueueItem[T, P])
    43  }
    44  
    45  // Update modifies the priority and value of an Item in the queue.
    46  func (pq *PriorityQueue[T, P]) Update(item *PriorityQueueItem[T, P], priority P) {
    47  	item.priority = priority
    48  	heap.Fix(pq, item.index)
    49  }
    50  
    51  func (pq *PriorityQueue[T, P]) Peek() *PriorityQueueItem[T, P] {
    52  	if len(*pq) > 0 {
    53  		return (*pq)[0]
    54  	}
    55  
    56  	return nil
    57  }
    58  
    59  func (pq *PriorityQueue[T, P]) Remove(item *PriorityQueueItem[T, P]) {
    60  	heap.Remove(pq, item.index)
    61  }
    62  
    63  func (pq PriorityQueue[T, P]) Len() int { return len(pq) }
    64  
    65  func (pq PriorityQueue[T, P]) Less(i, j int) bool {
    66  	return pq[i].priority < pq[j].priority
    67  }
    68  
    69  func (pq PriorityQueue[T, P]) Swap(i, j int) {
    70  	pq[i], pq[j] = pq[j], pq[i]
    71  	pq[i].index = i
    72  	pq[j].index = j
    73  }
    74  
    75  func (pq *PriorityQueue[T, P]) Push(x any) {
    76  	n := len(*pq)
    77  	item := x.(*PriorityQueueItem[T, P])
    78  	item.index = n
    79  	*pq = append(*pq, item)
    80  }
    81  
    82  func (pq *PriorityQueue[T, P]) Pop() any {
    83  	old := *pq
    84  	n := len(old)
    85  	item := old[n-1]
    86  	old[n-1] = nil
    87  	item.index = -1
    88  	*pq = old[0 : n-1]
    89  	return item
    90  }