refactor contact sort

This commit is contained in:
Alex Grintsvayg 2018-07-25 11:44:11 -04:00
parent c967af4a98
commit e642c110b8
3 changed files with 14 additions and 41 deletions

View file

@ -3,6 +3,7 @@ package dht
import ( import (
"bytes" "bytes"
"net" "net"
"sort"
"strconv" "strconv"
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
@ -115,15 +116,8 @@ func (c *Contact) UnmarshalBencode(b []byte) error {
return nil return nil
} }
type sortedContact struct { func sortByDistance(contacts []Contact, target bits.Bitmap) {
contact Contact sort.Slice(contacts, func(i, j int) bool {
xorDistanceToTarget bits.Bitmap return contacts[i].ID.Xor(target).Cmp(contacts[j].ID.Xor(target)) < 0
} })
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
} }

View file

@ -1,7 +1,6 @@
package dht package dht
import ( import (
"sort"
"sync" "sync"
"time" "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 // popFromShortlist pops the first contact off the shortlist and returns it
@ -345,17 +344,3 @@ func (cf *contactFinder) closest(contacts ...Contact) *Contact {
} }
return &closest 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
}
}

View file

@ -4,7 +4,6 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"net" "net"
"sort"
"strconv" "strconv"
"strings" "strings"
"sync" "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 // getClosest returns the closest `limit` contacts from the routing table
func (rt *routingTable) getClosest(target bits.Bitmap, limit int) []Contact { 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 var contacts []Contact
for _, sorted := range toSort { for _, b := range rt.buckets {
contacts = append(contacts, sorted.contact) contacts = append(contacts, b.Contacts()...)
if len(contacts) >= limit {
break
} }
sortByDistance(contacts, target)
if len(contacts) > limit {
contacts = contacts[:limit]
} }
return contacts return contacts
} }