github.com/MetalBlockchain/metalgo@v1.11.9/utils/math/averager_heap.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package math
     5  
     6  import (
     7  	"github.com/MetalBlockchain/metalgo/ids"
     8  	"github.com/MetalBlockchain/metalgo/utils/heap"
     9  )
    10  
    11  var _ AveragerHeap = (*averagerHeap)(nil)
    12  
    13  // TODO replace this interface with utils/heap
    14  // AveragerHeap maintains a heap of the averagers.
    15  type AveragerHeap interface {
    16  	// Add the average to the heap. If [nodeID] is already in the heap, the
    17  	// average will be replaced and the old average will be returned. If there
    18  	// was not an old average, false will be returned.
    19  	Add(nodeID ids.NodeID, averager Averager) (Averager, bool)
    20  	// Remove attempts to remove the average that was added with the provided
    21  	// [nodeID], if none is contained in the heap, [false] will be returned.
    22  	Remove(nodeID ids.NodeID) (Averager, bool)
    23  	// Pop attempts to remove the node with either the largest or smallest
    24  	// average, depending on if this is a max heap or a min heap, respectively.
    25  	Pop() (ids.NodeID, Averager, bool)
    26  	// Peek attempts to return the node with either the largest or smallest
    27  	// average, depending on if this is a max heap or a min heap, respectively.
    28  	Peek() (ids.NodeID, Averager, bool)
    29  	// Len returns the number of nodes that are currently in the heap.
    30  	Len() int
    31  }
    32  
    33  type averagerHeap struct {
    34  	heap heap.Map[ids.NodeID, Averager]
    35  }
    36  
    37  // NewMaxAveragerHeap returns a new empty max heap. The returned heap is not
    38  // thread safe.
    39  func NewMaxAveragerHeap() AveragerHeap {
    40  	return averagerHeap{
    41  		heap: heap.NewMap[ids.NodeID, Averager](func(a, b Averager) bool {
    42  			return a.Read() > b.Read()
    43  		}),
    44  	}
    45  }
    46  
    47  func (h averagerHeap) Add(nodeID ids.NodeID, averager Averager) (Averager, bool) {
    48  	return h.heap.Push(nodeID, averager)
    49  }
    50  
    51  func (h averagerHeap) Remove(nodeID ids.NodeID) (Averager, bool) {
    52  	return h.heap.Remove(nodeID)
    53  }
    54  
    55  func (h averagerHeap) Pop() (ids.NodeID, Averager, bool) {
    56  	return h.heap.Pop()
    57  }
    58  
    59  func (h averagerHeap) Peek() (ids.NodeID, Averager, bool) {
    60  	return h.heap.Peek()
    61  }
    62  
    63  func (h averagerHeap) Len() int {
    64  	return h.heap.Len()
    65  }