vendor/github.com/tendermint/tmlibs/common/heap.go

   1 package common
   2
   3 import (
   4         "container/heap"
   5 )
   6
   7 type Comparable interface {
   8         Less(o interface{}) bool
   9 }
  10
  11 //-----------------------------------------------------------------------------
  12
  13 /*
  14 Example usage:
  15         h := NewHeap()
  16
  17         h.Push(String("msg1"), 1)
  18         h.Push(String("msg3"), 3)
  19         h.Push(String("msg2"), 2)
  20
  21         fmt.Println(h.Pop())
  22         fmt.Println(h.Pop())
  23         fmt.Println(h.Pop())
  24 */
  25
  26 type Heap struct {
  27         pq priorityQueue
  28 }
  29
  30 func NewHeap() *Heap {
  31         return &Heap{pq: make([]*pqItem, 0)}
  32 }
  33
  34 func (h *Heap) Len() int64 {
  35         return int64(len(h.pq))
  36 }
  37
  38 func (h *Heap) Push(value interface{}, priority Comparable) {
  39         heap.Push(&h.pq, &pqItem{value: value, priority: priority})
  40 }
  41
  42 func (h *Heap) Peek() interface{} {
  43         if len(h.pq) == 0 {
  44                 return nil
  45         }
  46         return h.pq[0].value
  47 }
  48
  49 func (h *Heap) Update(value interface{}, priority Comparable) {
  50         h.pq.Update(h.pq[0], value, priority)
  51 }
  52
  53 func (h *Heap) Pop() interface{} {
  54         item := heap.Pop(&h.pq).(*pqItem)
  55         return item.value
  56 }
  57
  58 //-----------------------------------------------------------------------------
  59
  60 ///////////////////////
  61 // From: http://golang.org/pkg/container/heap/#example__priorityQueue
  62
  63 type pqItem struct {
  64         value    interface{}
  65         priority Comparable
  66         index    int
  67 }
  68
  69 type priorityQueue []*pqItem
  70
  71 func (pq priorityQueue) Len() int { return len(pq) }
  72
  73 func (pq priorityQueue) Less(i, j int) bool {
  74         return pq[i].priority.Less(pq[j].priority)
  75 }
  76
  77 func (pq priorityQueue) Swap(i, j int) {
  78         pq[i], pq[j] = pq[j], pq[i]
  79         pq[i].index = i
  80         pq[j].index = j
  81 }
  82
  83 func (pq *priorityQueue) Push(x interface{}) {
  84         n := len(*pq)
  85         item := x.(*pqItem)
  86         item.index = n
  87         *pq = append(*pq, item)
  88 }
  89
  90 func (pq *priorityQueue) Pop() interface{} {
  91         old := *pq
  92         n := len(old)
  93         item := old[n-1]
  94         item.index = -1 // for safety
  95         *pq = old[0 : n-1]
  96         return item
  97 }
  98
  99 func (pq *priorityQueue) Update(item *pqItem, value interface{}, priority Comparable) {
 100         item.value = value
 101         item.priority = priority
 102         heap.Fix(pq, item.index)
 103 }