refactor contact sort

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
+	})
 }

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
-	}
-}

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
 }