Add Heap Set (#2136)

Signed-off-by: Joshua Kim <20001595+joshua-kim@users.noreply.github.com>
Signed-off-by: Stephen Buttolph <stephen@avalabs.org>
Co-authored-by: Stephen Buttolph <stephen@avalabs.org>
Co-authored-by: Dhruba Basu <7675102+dhrubabasu@users.noreply.github.com>
Co-authored-by: dboehm-avalabs <david.boehm@avalabs.org>
Co-authored-by: David Boehm <91908103+dboehm-avalabs@users.noreply.github.com>
Co-authored-by: Alberto Benegiamo <alberto.benegiamo@gmail.com>

Author: 5 people

utils/heap/set.go

@@ -0,0 +1,48 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package heap
+
+// NewSet returns a heap without duplicates ordered by its values
+func NewSet[T comparable](less func(a, b T) bool) Set[T] {
+	return Set[T]{
+		set: NewMap[T, T](less),
+	}
+}
+
+type Set[T comparable] struct {
+	set Map[T, T]
+}
+
+// Push returns if the entry was added
+func (s Set[T]) Push(t T) bool {
+	_, hadValue := s.set.Push(t, t)
+	return !hadValue
+}
+
+func (s Set[T]) Pop() (T, bool) {
+	pop, _, ok := s.set.Pop()
+	return pop, ok
+}
+
+func (s Set[T]) Peek() (T, bool) {
+	peek, _, ok := s.set.Peek()
+	return peek, ok
+}
+
+func (s Set[T]) Len() int {
+	return s.set.Len()
+}
+
+func (s Set[T]) Remove(t T) bool {
+	_, existed := s.set.Remove(t)
+	return existed
+}
+
+func (s Set[T]) Fix(t T) {
+	s.set.Fix(t)
+}
+
+func (s Set[T]) Contains(t T) bool {
+	return s.set.Contains(t)
+}

utils/heap/set_test.go

@@ -0,0 +1,72 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package heap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSet(t *testing.T) {
+	tests := []struct {
+		name     string
+		setup    func(h Set[int])
+		expected []int
+	}{
+		{
+			name: "only push",
+			setup: func(h Set[int]) {
+				h.Push(1)
+				h.Push(2)
+				h.Push(3)
+			},
+			expected: []int{1, 2, 3},
+		},
+		{
+			name: "out of order pushes",
+			setup: func(h Set[int]) {
+				h.Push(1)
+				h.Push(5)
+				h.Push(2)
+				h.Push(4)
+				h.Push(3)
+			},
+			expected: []int{1, 2, 3, 4, 5},
+		},
+		{
+			name: "push and pop",
+			setup: func(h Set[int]) {
+				h.Push(1)
+				h.Push(5)
+				h.Push(2)
+				h.Push(4)
+				h.Push(3)
+				h.Pop()
+				h.Pop()
+				h.Pop()
+			},
+			expected: []int{4, 5},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			h := NewSet[int](func(a, b int) bool {
+				return a < b
+			})
+
+			tt.setup(h)
+
+			require.Equal(len(tt.expected), h.Len())
+			for _, expected := range tt.expected {
+				got, ok := h.Pop()
+				require.True(ok)
+				require.Equal(expected, got)
+			}
+		})
+	}
+}

x/sync/workheap.go

@@ -5,44 +5,38 @@ package sync
 
 import (
 	"bytes"
-	"container/heap"
 
+	"github.com/ava-labs/avalanchego/utils/heap"
 	"github.com/ava-labs/avalanchego/utils/math"
 	"github.com/ava-labs/avalanchego/utils/maybe"
 
 	"github.com/google/btree"
 )
 
-var _ heap.Interface = (*innerHeap)(nil)
-
-type heapItem struct {
-	workItem  *workItem
-	heapIndex int
-}
-
-type innerHeap []*heapItem
-
 // A priority queue of syncWorkItems.
 // Note that work item ranges never overlap.
 // Supports range merging and priority updating.
 // Not safe for concurrent use.
 type workHeap struct {
 	// Max heap of items by priority.
 	// i.e. heap.Pop returns highest priority item.
-	innerHeap innerHeap
+	innerHeap heap.Set[*workItem]
 	// The heap items sorted by range start.
 	// A Nothing start is considered to be the smallest.
-	sortedItems *btree.BTreeG[*heapItem]
+	sortedItems *btree.BTreeG[*workItem]
 	closed      bool
 }
 
 func newWorkHeap() *workHeap {
 	return &workHeap{
+		innerHeap: heap.NewSet[*workItem](func(a, b *workItem) bool {
+			return a.priority > b.priority
+		}),
 		sortedItems: btree.NewG(
 			2,
-			func(a, b *heapItem) bool {
-				aNothing := a.workItem.start.IsNothing()
-				bNothing := b.workItem.start.IsNothing()
+			func(a, b *workItem) bool {
+				aNothing := a.start.IsNothing()
+				bNothing := b.start.IsNothing()
 				if aNothing {
 					// [a] is Nothing, so if [b] is Nothing, they're equal.
 					// Otherwise, [b] is greater.
@@ -53,7 +47,7 @@ func newWorkHeap() *workHeap {
 					return false
 				}
 				// [a] and [b] both contain values. Compare the values.
-				return bytes.Compare(a.workItem.start.Value(), b.workItem.start.Value()) < 0
+				return bytes.Compare(a.start.Value(), b.start.Value()) < 0
 			},
 		),
 	}
@@ -70,10 +64,8 @@ func (wh *workHeap) Insert(item *workItem) {
 		return
 	}
 
-	wrappedItem := &heapItem{workItem: item}
-
-	heap.Push(&wh.innerHeap, wrappedItem)
-	wh.sortedItems.ReplaceOrInsert(wrappedItem)
+	wh.innerHeap.Push(item)
+	wh.sortedItems.ReplaceOrInsert(item)
 }
 
 // Pops and returns a work item from the heap.
@@ -82,9 +74,9 @@ func (wh *workHeap) GetWork() *workItem {
 	if wh.closed || wh.Len() == 0 {
 		return nil
 	}
-	item := heap.Pop(&wh.innerHeap).(*heapItem)
+	item, _ := wh.innerHeap.Pop()
 	wh.sortedItems.Delete(item)
-	return item.workItem
+	return item
 }
 
 // Insert the item into the heap, merging it with existing items
@@ -99,25 +91,23 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 		return
 	}
 
-	var mergedBefore, mergedAfter *heapItem
-	searchItem := &heapItem{
-		workItem: &workItem{
-			start: item.start,
-		},
+	var mergedBefore, mergedAfter *workItem
+	searchItem := &workItem{
+		start: item.start,
 	}
 
 	// Find the item with the greatest start range which is less than [item.start].
 	// Note that the iterator function will run at most once, since it always returns false.
 	wh.sortedItems.DescendLessOrEqual(
 		searchItem,
-		func(beforeItem *heapItem) bool {
-			if item.localRootID == beforeItem.workItem.localRootID &&
-				maybe.Equal(item.start, beforeItem.workItem.end, bytes.Equal) {
+		func(beforeItem *workItem) bool {
+			if item.localRootID == beforeItem.localRootID &&
+				maybe.Equal(item.start, beforeItem.end, bytes.Equal) {
 				// [beforeItem.start, beforeItem.end] and [item.start, item.end] are
 				// merged into [beforeItem.start, item.end]
-				beforeItem.workItem.end = item.end
-				beforeItem.workItem.priority = math.Max(item.priority, beforeItem.workItem.priority)
-				heap.Fix(&wh.innerHeap, beforeItem.heapIndex)
+				beforeItem.end = item.end
+				beforeItem.priority = math.Max(item.priority, beforeItem.priority)
+				wh.innerHeap.Fix(beforeItem)
 				mergedBefore = beforeItem
 			}
 			return false
@@ -127,14 +117,14 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 	// Note that the iterator function will run at most once, since it always returns false.
 	wh.sortedItems.AscendGreaterOrEqual(
 		searchItem,
-		func(afterItem *heapItem) bool {
-			if item.localRootID == afterItem.workItem.localRootID &&
-				maybe.Equal(item.end, afterItem.workItem.start, bytes.Equal) {
+		func(afterItem *workItem) bool {
+			if item.localRootID == afterItem.localRootID &&
+				maybe.Equal(item.end, afterItem.start, bytes.Equal) {
 				// [item.start, item.end] and [afterItem.start, afterItem.end] are merged into
 				// [item.start, afterItem.end].
-				afterItem.workItem.start = item.start
-				afterItem.workItem.priority = math.Max(item.priority, afterItem.workItem.priority)
-				heap.Fix(&wh.innerHeap, afterItem.heapIndex)
+				afterItem.start = item.start
+				afterItem.priority = math.Max(item.priority, afterItem.priority)
+				wh.innerHeap.Fix(afterItem)
 				mergedAfter = afterItem
 			}
 			return false
@@ -144,12 +134,12 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 	// we can combine the before item with the after item
 	if mergedBefore != nil && mergedAfter != nil {
 		// combine the two ranges
-		mergedBefore.workItem.end = mergedAfter.workItem.end
+		mergedBefore.end = mergedAfter.end
 		// remove the second range since it is now covered by the first
 		wh.remove(mergedAfter)
 		// update the priority
-		mergedBefore.workItem.priority = math.Max(mergedBefore.workItem.priority, mergedAfter.workItem.priority)
-		heap.Fix(&wh.innerHeap, mergedBefore.heapIndex)
+		mergedBefore.priority = math.Max(mergedBefore.priority, mergedAfter.priority)
+		wh.innerHeap.Fix(mergedBefore)
 	}
 
 	// nothing was merged, so add new item to the heap
@@ -160,43 +150,11 @@ func (wh *workHeap) MergeInsert(item *workItem) {
 }
 
 // Deletes [item] from the heap.
-func (wh *workHeap) remove(item *heapItem) {
-	heap.Remove(&wh.innerHeap, item.heapIndex)
-
+func (wh *workHeap) remove(item *workItem) {
+	wh.innerHeap.Remove(item)
 	wh.sortedItems.Delete(item)
 }
 
 func (wh *workHeap) Len() int {
 	return wh.innerHeap.Len()
 }
-
-// below this line are the implementations required for heap.Interface
-
-func (h innerHeap) Len() int {
-	return len(h)
-}
-
-func (h innerHeap) Less(i int, j int) bool {
-	return h[i].workItem.priority > h[j].workItem.priority
-}
-
-func (h innerHeap) Swap(i int, j int) {
-	h[i], h[j] = h[j], h[i]
-	h[i].heapIndex = i
-	h[j].heapIndex = j
-}
-
-func (h *innerHeap) Pop() interface{} {
-	old := *h
-	n := len(old)
-	item := old[n-1]
-	old[n-1] = nil
-	*h = old[0 : n-1]
-	return item
-}
-
-func (h *innerHeap) Push(x interface{}) {
-	item := x.(*heapItem)
-	item.heapIndex = len(*h)
-	*h = append(*h, item)
-}