Generic Heap

examples/tokenvm/controller/order_book.go

@@ -4,12 +4,11 @@
 package controller
 
 import (
-	"container/heap"
 	"sync"
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/hypersdk/crypto"
-	"github.com/ava-labs/hypersdk/examples/tokenvm/utils"
+	"github.com/ava-labs/hypersdk/heap"
 	"go.uber.org/zap"
 )
 
@@ -34,23 +33,23 @@ type OrderBook struct {
 	// TODO: consider capping the number of orders in each heap (need to ensure
 	// that doing so does not make it possible to send a bunch of small, spam
 	// orders to clear -> may need to set a min order limit to watch)
-	orders      map[string]*utils.Float64Heap[*Order]
+	orders      map[string]*heap.Heap[*Order, float64]
 	orderToPair map[ids.ID]string // needed to delete from [CloseOrder] actions
 	l           sync.RWMutex
 
 	trackAll bool
 }
 
 func NewOrderBook(c *Controller, trackedPairs []string) *OrderBook {
-	m := map[string]*utils.Float64Heap[*Order]{}
+	m := map[string]*heap.Heap[*Order, float64]{}
 	trackAll := false
 	if len(trackedPairs) == 1 && trackedPairs[0] == allPairs {
 		trackAll = true
 		c.inner.Logger().Info("tracking all order books")
 	} else {
 		for _, pair := range trackedPairs {
 			// We use a max heap so we return the best rates in order.
-			m[pair] = utils.NewFloat64Heap[*Order](initialPairCapacity, true)
+			m[pair] = heap.New[*Order, float64](initialPairCapacity, true)
 			c.inner.Logger().Info("tracking order book", zap.String("pair", pair))
 		}
 	}
@@ -71,10 +70,10 @@ func (o *OrderBook) Add(pair string, order *Order) {
 		return
 	case !ok && o.trackAll:
 		o.c.inner.Logger().Info("tracking order book", zap.String("pair", pair))
-		h = utils.NewFloat64Heap[*Order](initialPairCapacity, true)
+		h = heap.New[*Order, float64](initialPairCapacity, true)
 		o.orders[pair] = h
 	}
-	heap.Push(h, &utils.Float64Entry[*Order]{
+	h.Push(&heap.Entry[*Order, float64]{
 		ID:    order.ID,
 		Val:   float64(order.InTick) / float64(order.OutTick),
 		Item:  order,
@@ -96,12 +95,12 @@ func (o *OrderBook) Remove(id ids.ID) {
 		// This should never happen
 		return
 	}
-	entry, ok := h.GetID(id) // O(log 1)
+	entry, ok := h.Get(id) // O(log 1)
 	if !ok {
 		// This should never happen
 		return
 	}
-	heap.Remove(h, entry.Index) // O(log N)
+	h.Remove(entry.Index) // O(log N)
 }
 
 func (o *OrderBook) UpdateRemaining(id ids.ID, remaining uint64) {
@@ -116,7 +115,7 @@ func (o *OrderBook) UpdateRemaining(id ids.ID, remaining uint64) {
 		// This should never happen
 		return
 	}
-	entry, ok := h.GetID(id)
+	entry, ok := h.Get(id)
 	if !ok {
 		// This should never happen
 		return

heap/heap.go

@@ -0,0 +1,81 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package heap
+
+import (
+	"container/heap"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"golang.org/x/exp/constraints"
+)
+
+// Heap[I,V] is used to track objects of [I] by [Val].
+//
+// This data structure does not perform any synchronization and is not
+// safe to use concurrently without external locking.
+type Heap[I any, V constraints.Ordered] struct {
+	ih *innerHeap[I, V]
+}
+
+// New returns an instance of Heap[I,V]
+func New[I any, V constraints.Ordered](items int, isMinHeap bool) *Heap[I, V] {
+	return &Heap[I, V]{newInnerHeap[I, V](items, isMinHeap)}
+}
+
+// Len returns the number of items in ih.
+func (h *Heap[I, V]) Len() int { return h.ih.Len() }
+
+// Get returns the entry in th associated with [id], and a bool if [id] was
+// found in th.
+func (h *Heap[I, V]) Get(id ids.ID) (*Entry[I, V], bool) {
+	return h.ih.Get(id)
+}
+
+// Has returns whether [id] is found in th.
+func (h *Heap[I, V]) Has(id ids.ID) bool {
+	return h.ih.Has(id)
+}
+
+// Items returns all items in the heap in sorted order. You should not modify
+// the response.
+func (h *Heap[I, V]) Items() []*Entry[I, V] {
+	return h.ih.items
+}
+
+// Push can be called by external users instead of using `containers.heap`,
+// which makes using this heap less error-prone.
+func (h *Heap[I, V]) Push(e *Entry[I, V]) {
+	heap.Push(h.ih, e)
+}
+
+// Pop can be called by external users to remove an object from the heap at
+// a specific index instead of using `containers.heap`,
+// which makes using this heap less error-prone.
+func (h *Heap[I, V]) Pop() *Entry[I, V] {
+	if len(h.ih.items) == 0 {
+		return nil
+	}
+	return heap.Pop(h.ih).(*Entry[I, V])
+}
+
+// Remove can be called by external users to remove an object from the heap at
+// a specific index instead of using `containers.heap`,
+// which makes using this heap less error-prone.
+func (h *Heap[I, V]) Remove(index int) *Entry[I, V] {
+	if index >= len(h.ih.items) {
+		return nil
+	}
+	return heap.Remove(h.ih, index).(*Entry[I, V])
+}
+
+// First returns the first item in the heap. This is the smallest item in
+// a minHeap and the largest item in a maxHeap.
+//
+// If no items are in the heap, it will return nil.
+func (h *Heap[I, V]) First() *Entry[I, V] {
+	if len(h.ih.items) == 0 {
+		return nil
+	}
+	return h.ih.items[0]
+}