mirror of https://github.com/poanetwork/gecko.git
118 lines
2.9 KiB
Go
118 lines
2.9 KiB
Go
package avalanche
|
|
|
|
import (
|
|
"container/heap"
|
|
|
|
"github.com/ava-labs/gecko/ids"
|
|
"github.com/ava-labs/gecko/snow/consensus/avalanche"
|
|
)
|
|
|
|
// A vertexItem is a Vertex managed by the priority queue.
|
|
type vertexItem struct {
|
|
vertex avalanche.Vertex
|
|
index int // The index of the item in the heap.
|
|
}
|
|
|
|
// A priorityQueue implements heap.Interface and holds vertexItems.
|
|
type priorityQueue []*vertexItem
|
|
|
|
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].vertex.Status()
|
|
statusJ := pq[j].vertex.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
|
|
}
|
|
return pq[i].vertex.Height() > pq[j].vertex.Height()
|
|
}
|
|
|
|
func (pq priorityQueue) Swap(i, j int) {
|
|
pq[i], pq[j] = pq[j], pq[i]
|
|
pq[i].index = i
|
|
pq[j].index = j
|
|
}
|
|
|
|
// Push adds an item to this priority queue. x must have type *vertexItem
|
|
func (pq *priorityQueue) Push(x interface{}) {
|
|
n := len(*pq)
|
|
item := x.(*vertexItem)
|
|
item.index = n
|
|
*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
|
|
item.index = -1
|
|
*pq = old[0 : n-1]
|
|
return item
|
|
}
|
|
|
|
// vertexHeap defines the functionality of a heap of vertices
|
|
// with unique VertexIDs ordered by height
|
|
type vertexHeap interface {
|
|
Clear()
|
|
Push(avalanche.Vertex)
|
|
Pop() avalanche.Vertex // Requires that there be at least one element
|
|
Contains(avalanche.Vertex) bool
|
|
Len() int
|
|
}
|
|
|
|
type maxHeightVertexHeap struct {
|
|
heap *priorityQueue
|
|
elementIDs ids.Set
|
|
}
|
|
|
|
func newMaxVertexHeap() *maxHeightVertexHeap {
|
|
return &maxHeightVertexHeap{
|
|
heap: &priorityQueue{},
|
|
elementIDs: ids.Set{},
|
|
}
|
|
}
|
|
|
|
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)
|
|
item := &vertexItem{
|
|
vertex: vtx,
|
|
}
|
|
heap.Push(vh.heap, item)
|
|
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).(*vertexItem).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) }
|