storj/pkg/kademlia/queue.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
}