github.com/uber/kraken@v0.1.4/utils/heap/priority_queue.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package heap
    15  
    16  import (
    17  	"container/heap"
    18  	"errors"
    19  )
    20  
    21  // An Item is something we manage in a priority queue.
    22  type Item struct {
    23  	Value    interface{} // The value of the item; arbitrary.
    24  	Priority int         // The priority of the item in the queue (low value == high priority).
    25  }
    26  
    27  // PriorityQueue implements a heap returns Items with lowest priority first.
    28  type PriorityQueue struct {
    29  	q internalQueue
    30  }
    31  
    32  // NewPriorityQueue initializes a PriorityQueue with passed items.
    33  func NewPriorityQueue(items ...*Item) *PriorityQueue {
    34  	q := internalQueue(items)
    35  	heap.Init(&q)
    36  	return &PriorityQueue{q}
    37  }
    38  
    39  // Len returns the number of Items in the PriorityQueue.
    40  func (pq *PriorityQueue) Len() int { return len(pq.q) }
    41  
    42  // Push adds the Item to the PriorityQueue.
    43  func (pq *PriorityQueue) Push(item *Item) {
    44  	heap.Push(&pq.q, item)
    45  }
    46  
    47  // Pop removes and returns the lowest priority Item from the PriorityQueue.
    48  func (pq *PriorityQueue) Pop() (*Item, error) {
    49  	if len(pq.q) == 0 {
    50  		return nil, errors.New("queue empty")
    51  	}
    52  
    53  	return heap.Pop(&pq.q).(*Item), nil
    54  }
    55  
    56  // An internalQueue implements heap.Interface and holds Items.
    57  type internalQueue []*Item
    58  
    59  func (q internalQueue) Len() int { return len(q) }
    60  
    61  func (q internalQueue) Less(i, j int) bool {
    62  	return q[i].Priority < q[j].Priority
    63  }
    64  
    65  func (q internalQueue) Swap(i, j int) {
    66  	q[i], q[j] = q[j], q[i]
    67  }
    68  
    69  func (q *internalQueue) Push(x interface{}) {
    70  	item := x.(*Item)
    71  	*q = append(*q, item)
    72  }
    73  
    74  func (q *internalQueue) Pop() interface{} {
    75  	old := *q
    76  	n := len(old)
    77  	item := old[n-1]
    78  	*q = old[0 : n-1]
    79  	return item
    80  }