github.com/MetalBlockchain/metalgo@v1.11.9/utils/heap/queue.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package heap
     5  
     6  import (
     7  	"container/heap"
     8  
     9  	"github.com/MetalBlockchain/metalgo/utils"
    10  )
    11  
    12  var _ heap.Interface = (*queue[int])(nil)
    13  
    14  // NewQueue returns an empty heap. See QueueOf for more.
    15  func NewQueue[T any](less func(a, b T) bool) Queue[T] {
    16  	return QueueOf(less)
    17  }
    18  
    19  // QueueOf returns a heap containing entries ordered by less.
    20  func QueueOf[T any](less func(a, b T) bool, entries ...T) Queue[T] {
    21  	q := Queue[T]{
    22  		queue: &queue[T]{
    23  			entries: make([]T, len(entries)),
    24  			less:    less,
    25  		},
    26  	}
    27  
    28  	copy(q.queue.entries, entries)
    29  	heap.Init(q.queue)
    30  	return q
    31  }
    32  
    33  type Queue[T any] struct {
    34  	queue *queue[T]
    35  }
    36  
    37  func (q *Queue[T]) Len() int {
    38  	return len(q.queue.entries)
    39  }
    40  
    41  func (q *Queue[T]) Push(t T) {
    42  	heap.Push(q.queue, t)
    43  }
    44  
    45  func (q *Queue[T]) Pop() (T, bool) {
    46  	if q.Len() == 0 {
    47  		return utils.Zero[T](), false
    48  	}
    49  
    50  	return heap.Pop(q.queue).(T), true
    51  }
    52  
    53  func (q *Queue[T]) Peek() (T, bool) {
    54  	if q.Len() == 0 {
    55  		return utils.Zero[T](), false
    56  	}
    57  
    58  	return q.queue.entries[0], true
    59  }
    60  
    61  func (q *Queue[T]) Fix(i int) {
    62  	heap.Fix(q.queue, i)
    63  }
    64  
    65  type queue[T any] struct {
    66  	entries []T
    67  	less    func(a, b T) bool
    68  }
    69  
    70  func (q *queue[T]) Len() int {
    71  	return len(q.entries)
    72  }
    73  
    74  func (q *queue[T]) Less(i, j int) bool {
    75  	return q.less(q.entries[i], q.entries[j])
    76  }
    77  
    78  func (q *queue[T]) Swap(i, j int) {
    79  	q.entries[i], q.entries[j] = q.entries[j], q.entries[i]
    80  }
    81  
    82  func (q *queue[T]) Push(e any) {
    83  	q.entries = append(q.entries, e.(T))
    84  }
    85  
    86  func (q *queue[T]) Pop() any {
    87  	end := len(q.entries) - 1
    88  
    89  	popped := q.entries[end]
    90  	q.entries[end] = utils.Zero[T]()
    91  	q.entries = q.entries[:end]
    92  
    93  	return popped
    94  }