diff --git a/internal/eheap/eheap.go b/internal/eheap/eheap.go index 892d4f032a..79ecce3734 100644 --- a/internal/eheap/eheap.go +++ b/internal/eheap/eheap.go @@ -5,8 +5,7 @@ package eheap import ( "github.com/ava-labs/avalanchego/ids" - - "github.com/ava-labs/hypersdk/internal/heap" + "github.com/ava-labs/avalanchego/utils/heap" ) // Item is the interface that any item put in the heap must adheare to. @@ -22,39 +21,27 @@ type Item interface { // instead of grouping by expiry to support this feature, which makes it // less efficient). type ExpiryHeap[T Item] struct { - minHeap *heap.Heap[T, int64] + minHeap heap.Map[ids.ID, T] } // New returns an instance of ExpiryHeap with minHeap and maxHeap // containing [items]. -func New[T Item](items int) *ExpiryHeap[T] { +func New[T Item]() *ExpiryHeap[T] { return &ExpiryHeap[T]{ - minHeap: heap.New[T, int64](items, true), + minHeap: heap.NewMap[ids.ID, T](func(a, b T) bool { + return a.GetExpiry() < b.GetExpiry() + }), } } // Add pushes [item] to eh. func (eh *ExpiryHeap[T]) Add(item T) { - itemID := item.GetID() - poolLen := eh.minHeap.Len() - eh.minHeap.Push(&heap.Entry[T, int64]{ - ID: itemID, - Val: item.GetExpiry(), - Item: item, - Index: poolLen, - }) + eh.minHeap.Push(item.GetID(), item) } // Remove removes [id] from eh. If the id does not exist, Remove returns. func (eh *ExpiryHeap[T]) Remove(id ids.ID) (T, bool) { - minEntry, ok := eh.minHeap.Get(id) // O(1) - if !ok { - // This should never happen, as that would mean the heaps are out of - // sync. - return *new(T), false - } - eh.minHeap.Remove(minEntry.Index) // O(log N) - return minEntry.Item, true + return eh.minHeap.Remove(id) } // SetMin removes all elements in eh with a value less than [val]. Returns @@ -78,27 +65,19 @@ func (eh *ExpiryHeap[T]) SetMin(val int64) []T { // PeekMin returns the minimum value in eh. func (eh *ExpiryHeap[T]) PeekMin() (T, bool) { - first := eh.minHeap.First() - if first == nil { - return *new(T), false - } - return first.Item, true + _, v, ok := eh.minHeap.Peek() + return v, ok } // PopMin removes the minimum value in eh. func (eh *ExpiryHeap[T]) PopMin() (T, bool) { - first := eh.minHeap.First() - if first == nil { - return *new(T), false - } - item := first.Item - eh.Remove(item.GetID()) - return item, true + _, v, ok := eh.minHeap.Pop() + return v, ok } // Has returns if [item] is in eh. func (eh *ExpiryHeap[T]) Has(item ids.ID) bool { - return eh.minHeap.Has(item) + return eh.minHeap.Contains(item) } // Len returns the number of elements in eh. diff --git a/internal/eheap/eheap_test.go b/internal/eheap/eheap_test.go index 7d365545d2..7c648878eb 100644 --- a/internal/eheap/eheap_test.go +++ b/internal/eheap/eheap_test.go @@ -42,40 +42,40 @@ func GenerateTestItem(sponsor string, t int64) *TestItem { func TestExpiryHeapNew(t *testing.T) { // Creates empty min and max heaps require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() require.Zero(eheap.minHeap.Len(), "MinHeap not initialized correctly") } func TestExpiryHeapAdd(t *testing.T) { // Adds to the mempool. require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() item := GenerateTestItem("sponsor", 1) eheap.Add(item) - require.Equal(1, eheap.minHeap.Len(), "MinHeap not pushed correctly") - require.True(eheap.minHeap.Has(item.GetID()), "MinHeap does not have ID") + require.Equal(1, eheap.Len(), "MinHeap not pushed correctly") + require.True(eheap.Has(item.GetID()), "MinHeap does not have ID") } func TestExpiryHeapRemove(t *testing.T) { // Removes from the mempool. require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() item := GenerateTestItem("sponsor", 1) // Add first eheap.Add(item) require.Equal(1, eheap.minHeap.Len(), "MinHeap not pushed correctly") - require.True(eheap.minHeap.Has(item.GetID()), "MinHeap does not have ID") + require.True(eheap.Has(item.GetID()), "MinHeap does not have ID") // Remove eheap.Remove(item.GetID()) require.Zero(eheap.minHeap.Len(), "MinHeap not removed") - require.False(eheap.minHeap.Has(item.GetID()), "MinHeap still has ID") + require.False(eheap.Has(item.GetID()), "MinHeap still has ID") } func TestExpiryHeapRemoveEmpty(t *testing.T) { // Try to remove a non existing entry. // Removes from the mempool. require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() item := GenerateTestItem("sponsor", 1) // Require this returns eheap.Remove(item.GetID()) @@ -85,7 +85,7 @@ func TestExpiryHeapRemoveEmpty(t *testing.T) { func TestSetMin(t *testing.T) { require := require.New(t) sponsor := "sponsor" - eheap := New[*TestItem](0) + eheap := New[*TestItem]() for i := int64(0); i <= 9; i++ { item := GenerateTestItem(sponsor, i) eheap.Add(item) @@ -109,7 +109,7 @@ func TestSetMin(t *testing.T) { func TestSetMinRemovesAll(t *testing.T) { require := require.New(t) sponsor := "sponsor" - eheap := New[*TestItem](0) + eheap := New[*TestItem]() var items []*TestItem for i := int64(0); i <= 4; i++ { item := GenerateTestItem(sponsor, i) @@ -126,7 +126,7 @@ func TestSetMinRemovesAll(t *testing.T) { func TestPeekMin(t *testing.T) { require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() itemMin := GenerateTestItem(testSponsor, 1) itemMed := GenerateTestItem(testSponsor, 2) @@ -157,7 +157,7 @@ func TestPeekMin(t *testing.T) { func TestPopMin(t *testing.T) { require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() itemMin := GenerateTestItem(testSponsor, 1) itemMed := GenerateTestItem(testSponsor, 2) @@ -183,7 +183,7 @@ func TestPopMin(t *testing.T) { func TestHas(t *testing.T) { require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() item := GenerateTestItem(testSponsor, 1) require.False(eheap.Has(item.GetID()), "Found an item that was not added.") eheap.Add(item) @@ -193,7 +193,7 @@ func TestHas(t *testing.T) { func TestLen(t *testing.T) { require := require.New(t) - eheap := New[*TestItem](0) + eheap := New[*TestItem]() for i := int64(0); i <= 4; i++ { item := GenerateTestItem(testSponsor, i) eheap.Add(item) diff --git a/internal/mempool/mempool.go b/internal/mempool/mempool.go index 062b034c3e..da0528a3ca 100644 --- a/internal/mempool/mempool.go +++ b/internal/mempool/mempool.go @@ -66,7 +66,7 @@ func New[T Item]( maxSponsorSize: maxSponsorSize, queue: &list.List[T]{}, - eh: eheap.New[*list.Element[T]](min(maxSize, maxPrealloc)), + eh: eheap.New[*list.Element[T]](), owned: map[codec.Address]int{}, }