diff --git a/utils/heap/queue.go b/utils/heap/queue.go index 4d05158265c5..e5c446afcd9c 100644 --- a/utils/heap/queue.go +++ b/utils/heap/queue.go @@ -10,14 +10,22 @@ import ( var _ stdheap.Interface = (*queue[int])(nil) -// NewQueue returns an empty min-heap ordered by less. See OfQueue for more. +// NewQueue returns an empty min-heap ordered by less. See QueueOf for more. func NewQueue[T any](less func(a, b T) bool) Queue[T] { - return Queue[T]{ + return QueueOf(less) +} + +func QueueOf[T any](less func(a, b T) bool, entries ...T) Queue[T] { + q := Queue[T]{ queue: &queue[T]{ entries: make([]T, 0), less: less, }, } + + for _, e := range entries { + q.Push(e) + } } type Queue[T any] struct { diff --git a/vms/platformvm/state/merged_iterator.go b/vms/platformvm/state/merged_iterator.go index 2be90a1a3106..08a7e34811b2 100644 --- a/vms/platformvm/state/merged_iterator.go +++ b/vms/platformvm/state/merged_iterator.go @@ -3,18 +3,15 @@ package state -import "container/heap" +import "github.com/ava-labs/avalanchego/utils/heap" -var ( - _ StakerIterator = (*mergedIterator)(nil) - _ heap.Interface = (*mergedIterator)(nil) -) +var _ StakerIterator = (*mergedIterator)(nil) type mergedIterator struct { initialized bool // heap only contains iterators that have been initialized and are not // exhausted. - heap []StakerIterator + heap heap.Queue[StakerIterator] } // Returns an iterator that returns all of the elements of [stakers] in order. @@ -36,15 +33,19 @@ func NewMergedIterator(stakers ...StakerIterator) StakerIterator { } it := &mergedIterator{ - heap: stakers, + heap: heap.OfQueue[StakerIterator]( + func(a, b StakerIterator) bool { + return a.Value().Less(b.Value()) + }, + stakers..., + ), } - heap.Init(it) return it } func (it *mergedIterator) Next() bool { - if len(it.heap) == 0 { + if it.heap.Len() == 0 { return false } @@ -57,54 +58,31 @@ func (it *mergedIterator) Next() bool { } // Update the heap root. - current := it.heap[0] + current, _ := it.heap.Peek() if current.Next() { // Calling Next() above modifies [current] so we fix the heap. - heap.Fix(it, 0) + it.heap.Fix(0) return true } // The old root is exhausted. Remove it from the heap. current.Release() - heap.Pop(it) - return len(it.heap) > 0 + it.heap.Pop() + return it.heap.Len() > 0 } func (it *mergedIterator) Value() *Staker { - return it.heap[0].Value() + peek, _ := it.heap.Peek() + return peek.Value() } -// When Release() returns, Release() has been called on each element of -// [stakers]. func (it *mergedIterator) Release() { - for _, it := range it.heap { - it.Release() + for it.heap.Len() > 0 { + removed, _ := it.heap.Pop() + removed.Release() } - it.heap = nil } -// Returns the number of sub-iterators in [it]. func (it *mergedIterator) Len() int { - return len(it.heap) -} - -func (it *mergedIterator) Less(i, j int) bool { - return it.heap[i].Value().Less(it.heap[j].Value()) -} - -func (it *mergedIterator) Swap(i, j int) { - it.heap[j], it.heap[i] = it.heap[i], it.heap[j] -} - -// Push is never actually used - but we need it to implement heap.Interface. -func (it *mergedIterator) Push(value interface{}) { - it.heap = append(it.heap, value.(StakerIterator)) -} - -func (it *mergedIterator) Pop() interface{} { - newLength := len(it.heap) - 1 - value := it.heap[newLength] - it.heap[newLength] = nil - it.heap = it.heap[:newLength] - return value + return it.heap.Len() }