status-go/vendor/github.com/anacrolix/dht/v2/table.go

135 lines
2.8 KiB
Go
Raw Normal View History

2022-03-10 09:44:48 +00:00
package dht
import (
"errors"
"fmt"
"github.com/anacrolix/dht/v2/int160"
)
// Node table, with indexes on distance from root ID to bucket, and node addr.
type table struct {
rootID int160.T
k int
buckets [160]bucket
addrs map[string]map[int160.T]struct{}
}
func (tbl *table) K() int {
return tbl.k
}
func (tbl *table) randomIdForBucket(bucketIndex int) int160.T {
randomId := randomIdInBucket(tbl.rootID, bucketIndex)
if randomIdBucketIndex := tbl.bucketIndex(randomId); randomIdBucketIndex != bucketIndex {
panic(fmt.Sprintf("bucket index for random id %v == %v not %v", randomId, randomIdBucketIndex, bucketIndex))
}
return randomId
}
func (tbl *table) addrNodes(addr Addr) []*node {
a := tbl.addrs[addr.String()]
ret := make([]*node, 0, len(a))
for id := range a {
ret = append(ret, tbl.getNode(addr, id))
}
return ret
}
func (tbl *table) dropNode(n *node) {
as := n.Addr.String()
if _, ok := tbl.addrs[as][n.Id]; !ok {
panic("missing id for addr")
}
delete(tbl.addrs[as], n.Id)
if len(tbl.addrs[as]) == 0 {
delete(tbl.addrs, as)
}
b := tbl.bucketForID(n.Id)
if _, ok := b.nodes[n]; !ok {
panic("expected node in bucket")
}
delete(b.nodes, n)
}
func (tbl *table) bucketForID(id int160.T) *bucket {
return &tbl.buckets[tbl.bucketIndex(id)]
}
func (tbl *table) numNodes() (num int) {
for _, b := range tbl.buckets {
num += b.Len()
}
return
}
func (tbl *table) bucketIndex(id int160.T) int {
if id == tbl.rootID {
panic("nobody puts the root ID in a bucket")
}
var a int160.T
a.Xor(&tbl.rootID, &id)
index := 160 - a.BitLen()
return index
}
func (tbl *table) forNodes(f func(*node) bool) bool {
for _, b := range tbl.buckets {
if !b.EachNode(f) {
return false
}
}
return true
}
func (tbl *table) getNode(addr Addr, id int160.T) *node {
if id == tbl.rootID {
return nil
}
return tbl.buckets[tbl.bucketIndex(id)].GetNode(addr, id)
}
func (tbl *table) closestNodes(k int, target int160.T, filter func(*node) bool) (ret []*node) {
for bi := func() int {
if target == tbl.rootID {
return len(tbl.buckets) - 1
} else {
return tbl.bucketIndex(target)
}
}(); bi >= 0 && len(ret) < k; bi-- {
for n := range tbl.buckets[bi].nodes {
if filter(n) {
ret = append(ret, n)
}
}
}
// TODO: Keep only the closest.
if len(ret) > k {
ret = ret[:k]
}
return
}
func (tbl *table) addNode(n *node) error {
if n.Id == tbl.rootID {
return errors.New("is root id")
}
b := &tbl.buckets[tbl.bucketIndex(n.Id)]
if b.GetNode(n.Addr, n.Id) != nil {
return errors.New("already present")
}
if b.Len() >= tbl.k {
return errors.New("bucket is full")
}
b.AddNode(n, tbl.k)
if tbl.addrs == nil {
tbl.addrs = make(map[string]map[int160.T]struct{}, 160*tbl.k)
}
as := n.Addr.String()
if tbl.addrs[as] == nil {
tbl.addrs[as] = make(map[int160.T]struct{}, 1)
}
tbl.addrs[as][n.Id] = struct{}{}
return nil
}