-
Notifications
You must be signed in to change notification settings - Fork 649
/
heap.go
135 lines (110 loc) · 3.21 KB
/
heap.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package vertex
import (
"container/heap"
"github.com/ava-labs/avalanchego/ids"
"github.com/ava-labs/avalanchego/snow/consensus/avalanche"
"github.com/ava-labs/avalanchego/utils/set"
)
var (
_ Heap = (*maxHeightVertexHeap)(nil)
_ heap.Interface = (*priorityQueue)(nil)
)
type priorityQueue []avalanche.Vertex
func (pq priorityQueue) Len() int {
return len(pq)
}
// Returns true if the vertex at index i has greater height than the vertex at
// index j.
func (pq priorityQueue) Less(i, j int) bool {
statusI := pq[i].Status()
statusJ := pq[j].Status()
// Put unknown vertices at the front of the heap to ensure once we have made
// it below a certain height in DAG traversal we do not need to reset
if !statusI.Fetched() {
return true
}
if !statusJ.Fetched() {
return false
}
// Treat errors on retrieving the height as if the vertex is not fetched
heightI, errI := pq[i].Height()
if errI != nil {
return true
}
heightJ, errJ := pq[j].Height()
if errJ != nil {
return false
}
return heightI > heightJ
}
func (pq priorityQueue) Swap(i, j int) {
pq[i], pq[j] = pq[j], pq[i]
}
// Push adds an item to this priority queue. x must have type *vertexItem
func (pq *priorityQueue) Push(x interface{}) {
item := x.(avalanche.Vertex)
*pq = append(*pq, item)
}
// Pop returns the last item in this priorityQueue
func (pq *priorityQueue) Pop() interface{} {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil
*pq = old[0 : n-1]
return item
}
// Heap defines the functionality of a heap of vertices with unique VertexIDs
// ordered by height
type Heap interface {
// Empty the heap.
Clear()
// Add the provided vertex to the heap. Vertices are de-duplicated, returns
// true if the vertex was added, false if it was dropped.
Push(avalanche.Vertex) bool
// Remove the top vertex. Assumes that there is at least one element.
Pop() avalanche.Vertex
// Returns if a vertex with the provided ID is currently in the heap.
Contains(ids.ID) bool
// Returns the number of vertices in the heap.
Len() int
}
// NewHeap returns an empty Heap
func NewHeap() Heap {
return &maxHeightVertexHeap{}
}
type maxHeightVertexHeap struct {
heap priorityQueue
elementIDs set.Set[ids.ID]
}
func (vh *maxHeightVertexHeap) Clear() {
vh.heap = priorityQueue{}
vh.elementIDs.Clear()
}
// Push adds an element to this heap. Returns true if the element was added.
// Returns false if it was already in the heap.
func (vh *maxHeightVertexHeap) Push(vtx avalanche.Vertex) bool {
vtxID := vtx.ID()
if vh.elementIDs.Contains(vtxID) {
return false
}
vh.elementIDs.Add(vtxID)
heap.Push(&vh.heap, vtx)
return true
}
// If there are any vertices in this heap with status Unknown, removes one such
// vertex and returns it. Otherwise, removes and returns the vertex in this heap
// with the greatest height.
func (vh *maxHeightVertexHeap) Pop() avalanche.Vertex {
vtx := heap.Pop(&vh.heap).(avalanche.Vertex)
vh.elementIDs.Remove(vtx.ID())
return vtx
}
func (vh *maxHeightVertexHeap) Len() int {
return vh.heap.Len()
}
func (vh *maxHeightVertexHeap) Contains(vtxID ids.ID) bool {
return vh.elementIDs.Contains(vtxID)
}