Alternative changes to make set heap work (#2146)

ava-labs · Oct 6, 2023 · 5d0cfd0 · 5d0cfd0
1 parent d900206
commit 5d0cfd0
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/utils/heap/set.go b/utils/heap/set.go
@@ -14,10 +14,9 @@ type Set[T comparable] struct {
 	set Map[T, T]
 }
 
-// Push returns if a value was overwritten
-func (s Set[T]) Push(t T) bool {
-	_, ok := s.set.Push(t, t)
-	return ok
+// Push returns the evicted previous value if present
+func (s Set[T]) Push(t T) (T, bool) {
+	return s.set.Push(t, t)
 }
 
 func (s Set[T]) Pop() (T, bool) {
@@ -34,15 +33,20 @@ func (s Set[T]) Len() int {
 	return s.set.Len()
 }
 
-func (s Set[T]) Remove(i int) T {
-	remove, _ := s.set.Remove(i)
-	return remove
+func (s Set[T]) Remove(t T) (T, bool) {
+	remove, _, existed := s.set.Remove(t)
+	return remove, existed
 }
 
-func (s Set[T]) Fix(i int) {
-	s.set.Fix(i)
+func (s Set[T]) Fix(t T) {
+	s.set.Fix(t)
 }
 
-func (s Set[T]) Index() map[T]int {
-	return s.set.queue.index
+func (s Set[T]) Contains(t T) bool {
+	return s.set.Contains(t)
+}
+
+func (s Set[T]) Get(t T) (T, bool) {
+	key, _, ok := s.set.Get(t)
+	return key, ok
 }
diff --git a/x/sync/workheap.go b/x/sync/workheap.go
@@ -107,7 +107,7 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 				// merged into [beforeItem.start, item.end]
 				beforeItem.end = item.end
 				beforeItem.priority = math.Max(item.priority, beforeItem.priority)
-				wh.innerHeap.Fix(wh.innerHeap.Index()[beforeItem])
+				wh.innerHeap.Fix(beforeItem)
 				mergedBefore = beforeItem
 			}
 			return false
@@ -124,7 +124,7 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 				// [item.start, afterItem.end].
 				afterItem.start = item.start
 				afterItem.priority = math.Max(item.priority, afterItem.priority)
-				wh.innerHeap.Fix(wh.innerHeap.Index()[afterItem])
+				wh.innerHeap.Fix(afterItem)
 				mergedAfter = afterItem
 			}
 			return false
@@ -139,7 +139,7 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 		wh.remove(mergedAfter)
 		// update the priority
 		mergedBefore.priority = math.Max(mergedBefore.priority, mergedAfter.priority)
-		wh.innerHeap.Fix(wh.innerHeap.Index()[mergedBefore])
+		wh.innerHeap.Fix(mergedBefore)
 	}
 
 	// nothing was merged, so add new item to the heap
@@ -151,7 +151,7 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 
 // Deletes [item] from the heap.
 func (wh *workHeap) remove(item *workItem) {
-	wh.innerHeap.Remove(wh.innerHeap.Index()[item])
+	wh.innerHeap.Remove(item)
 	wh.sortedItems.Delete(item)
 }