842ebc9546
* PriorityQueue refactor * changed closest to return big.Int * test smaller queue, return nil on empty
106 lines
2.6 KiB
Go
106 lines
2.6 KiB
Go
// Copyright (C) 2018 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
package kademlia
|
|
|
|
import (
|
|
"container/heap"
|
|
"math/big"
|
|
|
|
"storj.io/storj/pkg/dht"
|
|
"storj.io/storj/pkg/pb"
|
|
)
|
|
|
|
//XorQueue is a priority queue where the priority is key XOR distance
|
|
type XorQueue struct {
|
|
maxLen int
|
|
items items
|
|
}
|
|
|
|
//NewXorQueue returns a items with priority based on XOR from targetBytes
|
|
func NewXorQueue(size int) *XorQueue {
|
|
return &XorQueue{items: make(items, 0, size), maxLen: size}
|
|
}
|
|
|
|
//Insert adds Node onto the queue
|
|
func (x *XorQueue) Insert(target dht.NodeID, nodes []*pb.Node) {
|
|
targetBytes := new(big.Int).SetBytes(target.Bytes())
|
|
//insert new nodes
|
|
for _, node := range nodes {
|
|
heap.Push(&x.items, &item{
|
|
value: node,
|
|
priority: new(big.Int).Xor(targetBytes, new(big.Int).SetBytes([]byte(node.GetId()))),
|
|
})
|
|
}
|
|
//resize down if we grew too big
|
|
if x.items.Len() > x.maxLen {
|
|
olditems := x.items
|
|
x.items = items{}
|
|
for i := 0; i < x.maxLen && len(olditems) > 0; i++ {
|
|
item := heap.Pop(&olditems)
|
|
heap.Push(&x.items, item)
|
|
}
|
|
heap.Init(&x.items)
|
|
}
|
|
}
|
|
|
|
//Closest removed the closest priority node from the queue
|
|
func (x *XorQueue) Closest() (*pb.Node, big.Int) {
|
|
if x.Len() == 0 {
|
|
return nil, big.Int{}
|
|
}
|
|
item := *(heap.Pop(&x.items).(*item))
|
|
return item.value, *item.priority
|
|
}
|
|
|
|
//Len returns the number of items in the queue
|
|
func (x *XorQueue) Len() int {
|
|
return x.items.Len()
|
|
}
|
|
|
|
// An item is something we manage in a priority queue.
|
|
type item struct {
|
|
value *pb.Node // The value of the item; arbitrary.
|
|
priority *big.Int // The priority of the item in the queue.
|
|
// The index is needed by update and is maintained by the heap.Interface methods.
|
|
index int // The index of the item in the heap.
|
|
}
|
|
|
|
// A items implements heap.Interface and holds items.
|
|
type items []*item
|
|
|
|
// Len returns the length of the priority queue
|
|
func (items items) Len() int { return len(items) }
|
|
|
|
// Less does what you would think
|
|
func (items items) Less(i, j int) bool {
|
|
// this sorts the nodes where the node popped has the closest location
|
|
return items[i].priority.Cmp(items[j].priority) < 0
|
|
}
|
|
|
|
// Swap swaps two ints
|
|
func (items items) Swap(i, j int) {
|
|
items[i], items[j] = items[j], items[i]
|
|
items[i].index = i
|
|
items[j].index = j
|
|
}
|
|
|
|
// Push adds an item to the top of the queue
|
|
// must call heap.fix to resort
|
|
func (items *items) Push(x interface{}) {
|
|
n := len(*items)
|
|
item := x.(*item)
|
|
item.index = n
|
|
*items = append(*items, item)
|
|
}
|
|
|
|
// Pop returns the item with the lowest priority
|
|
func (items *items) Pop() interface{} {
|
|
old := *items
|
|
n := len(old)
|
|
item := old[n-1]
|
|
item.index = -1 // for safety
|
|
*items = old[0 : n-1]
|
|
return item
|
|
}
|