diff --git a/dht/contact.go b/dht/contact.go index 64f205b..027b258 100644 --- a/dht/contact.go +++ b/dht/contact.go @@ -3,6 +3,7 @@ package dht import ( "bytes" "net" + "sort" "strconv" "github.com/lbryio/lbry.go/errors" @@ -115,15 +116,8 @@ func (c *Contact) UnmarshalBencode(b []byte) error { return nil } -type sortedContact struct { - contact Contact - xorDistanceToTarget bits.Bitmap -} - -type byXorDistance []sortedContact - -func (a byXorDistance) Len() int { return len(a) } -func (a byXorDistance) Swap(i, j int) { a[i], a[j] = a[j], a[i] } -func (a byXorDistance) Less(i, j int) bool { - return a[i].xorDistanceToTarget.Cmp(a[j].xorDistanceToTarget) < 0 +func sortByDistance(contacts []Contact, target bits.Bitmap) { + sort.Slice(contacts, func(i, j int) bool { + return contacts[i].ID.Xor(target).Cmp(contacts[j].ID.Xor(target)) < 0 + }) } diff --git a/dht/node_finder.go b/dht/node_finder.go index 7375bd9..60b2098 100644 --- a/dht/node_finder.go +++ b/dht/node_finder.go @@ -1,7 +1,6 @@ package dht import ( - "sort" "sync" "time" @@ -268,7 +267,7 @@ func (cf *contactFinder) appendNewToShortlist(contacts []Contact) { } } - sortInPlace(cf.shortlist, cf.target) + sortByDistance(cf.shortlist, cf.target) } // popFromShortlist pops the first contact off the shortlist and returns it @@ -345,17 +344,3 @@ func (cf *contactFinder) closest(contacts ...Contact) *Contact { } return &closest } - -func sortInPlace(contacts []Contact, target bits.Bitmap) { - toSort := make([]sortedContact, len(contacts)) - - for i, n := range contacts { - toSort[i] = sortedContact{n, n.ID.Xor(target)} - } - - sort.Sort(byXorDistance(toSort)) - - for i, c := range toSort { - contacts[i] = c.contact - } -} diff --git a/dht/routing_table.go b/dht/routing_table.go index 7b53d9a..30fc906 100644 --- a/dht/routing_table.go +++ b/dht/routing_table.go @@ -4,7 +4,6 @@ import ( "encoding/json" "fmt" "net" - "sort" "strconv" "strings" "sync" @@ -291,21 +290,16 @@ func (rt *routingTable) GetClosest(target bits.Bitmap, limit int) []Contact { // getClosest returns the closest `limit` contacts from the routing table func (rt *routingTable) getClosest(target bits.Bitmap, limit int) []Contact { - var toSort []sortedContact - for _, b := range rt.buckets { - for _, c := range b.Contacts() { - toSort = append(toSort, sortedContact{c, c.ID.Xor(target)}) - } - } - sort.Sort(byXorDistance(toSort)) - var contacts []Contact - for _, sorted := range toSort { - contacts = append(contacts, sorted.contact) - if len(contacts) >= limit { - break - } + for _, b := range rt.buckets { + contacts = append(contacts, b.Contacts()...) } + + sortByDistance(contacts, target) + if len(contacts) > limit { + contacts = contacts[:limit] + } + return contacts }