Expose AddrMaybeIdsByDistance

Author: anacrolix

containers/addr-maybe-ids-by-distance.go

@@ -0,0 +1,15 @@
+package containers
+
+import (
+	"github.com/anacrolix/dht/v2/int160"
+	"github.com/anacrolix/dht/v2/types"
+	"github.com/anacrolix/stm/stmutil"
+)
+
+type addrMaybeId = types.AddrMaybeId
+
+func NewAddrMaybeIdsByDistance(target int160.T) stmutil.Settish {
+	return stmutil.NewSortedSet(func(l, r interface{}) bool {
+		return l.(addrMaybeId).CloserThan(r.(addrMaybeId), target)
+	})
+}

containers/addr-maybe-ids-by-distance_test.go

@@ -0,0 +1,41 @@
+package containers
+
+import (
+	"testing"
+
+	"github.com/anacrolix/dht/v2/int160"
+	"github.com/anacrolix/dht/v2/internal/testutil"
+	"github.com/anacrolix/missinggo/v2/iter"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNodesByDistance(t *testing.T) {
+	a := NewAddrMaybeIdsByDistance(int160.T{})
+	push := func(i int) {
+		a = a.Add(testutil.SampleAddrMaybeIds[i])
+	}
+	push(4)
+	push(2)
+	push(0)
+	push(3)
+	push(0)
+	push(1)
+	pop := func(is ...int) {
+		first, ok := iter.First(a.Iter)
+		assert.True(t, ok)
+		assert.Contains(t, func() (ret []addrMaybeId) {
+			for _, i := range is {
+				ret = append(ret, testutil.SampleAddrMaybeIds[i])
+			}
+			return
+		}(), first)
+		a = a.Delete(first)
+	}
+	pop(1)
+	pop(2)
+	pop(3)
+	pop(0, 4)
+	pop(0, 4)
+	//pop(0, 4)
+	assert.EqualValues(t, 0, a.Len())
+}

internal/testutil/testutil.go

@@ -0,0 +1,26 @@
+package testutil
+
+import (
+	"github.com/anacrolix/dht/v2/int160"
+	"github.com/anacrolix/dht/v2/krpc"
+	"github.com/anacrolix/dht/v2/types"
+)
+
+func Int160WithBitSet(bit int) *int160.T {
+	var i int160.T
+	i.SetBit(7+bit*8, true)
+	return &i
+}
+
+type addrMaybeId = types.AddrMaybeId
+
+var SampleAddrMaybeIds = []addrMaybeId{
+	addrMaybeId{},
+	addrMaybeId{Id: new(int160.T)},
+	addrMaybeId{Id: Int160WithBitSet(13)},
+	addrMaybeId{Id: Int160WithBitSet(12)},
+	addrMaybeId{Addr: krpc.NodeAddr{Port: 1}},
+	addrMaybeId{
+		Id:   Int160WithBitSet(14),
+		Addr: krpc.NodeAddr{Port: 1}},
+}

misc.go

@@ -3,11 +3,10 @@ package dht
 import (
 	"net"
 
-	"github.com/anacrolix/missinggo/v2/iter"
-	"github.com/anacrolix/stm/stmutil"
-
 	"github.com/anacrolix/dht/v2/int160"
+	"github.com/anacrolix/dht/v2/krpc"
 	"github.com/anacrolix/dht/v2/types"
+	"github.com/anacrolix/missinggo/v2/iter"
 )
 
 func mustListen(addr string) net.PacketConn {
@@ -27,14 +26,8 @@ func addrResolver(addr string) func() ([]Addr, error) {
 
 type addrMaybeId = types.AddrMaybeId
 
-func nodesByDistance(target int160.T) stmutil.Settish {
-	return stmutil.NewSortedSet(func(l, r interface{}) bool {
-		return l.(addrMaybeId).CloserThan(r.(addrMaybeId), target)
-	})
-}
-
 func randomIdInBucket(rootId int160.T, bucketIndex int) int160.T {
-	id := int160.FromByteArray(RandomNodeID())
+	id := int160.FromByteArray(krpc.RandomNodeID())
 	for i := range iter.N(bucketIndex) {
 		id.SetBit(i, rootId.GetBit(i))
 	}

misc_test.go

@@ -1,59 +1 @@
 package dht
-
-import (
-	"testing"
-
-	"github.com/anacrolix/missinggo/v2/iter"
-	"github.com/stretchr/testify/assert"
-
-	"github.com/anacrolix/dht/v2/int160"
-	"github.com/anacrolix/dht/v2/krpc"
-)
-
-func int160WithBitSet(bit int) *int160.T {
-	var i int160.T
-	i.SetBit(7+bit*8, true)
-	return &i
-}
-
-var sampleAddrMaybeIds = []addrMaybeId{
-	addrMaybeId{},
-	addrMaybeId{Id: new(int160.T)},
-	addrMaybeId{Id: int160WithBitSet(13)},
-	addrMaybeId{Id: int160WithBitSet(12)},
-	addrMaybeId{Addr: krpc.NodeAddr{Port: 1}},
-	addrMaybeId{
-		Id:   int160WithBitSet(14),
-		Addr: krpc.NodeAddr{Port: 1}},
-}
-
-func TestNodesByDistance(t *testing.T) {
-	a := nodesByDistance(int160.T{})
-	push := func(i int) {
-		a = a.Add(sampleAddrMaybeIds[i])
-	}
-	push(4)
-	push(2)
-	push(0)
-	push(3)
-	push(0)
-	push(1)
-	pop := func(is ...int) {
-		first, ok := iter.First(a.Iter)
-		assert.True(t, ok)
-		assert.Contains(t, func() (ret []addrMaybeId) {
-			for _, i := range is {
-				ret = append(ret, sampleAddrMaybeIds[i])
-			}
-			return
-		}(), first)
-		a = a.Delete(first)
-	}
-	pop(1)
-	pop(2)
-	pop(3)
-	pop(0, 4)
-	pop(0, 4)
-	//pop(0, 4)
-	assert.EqualValues(t, 0, a.Len())
-}

stm-traversal.go

@@ -7,6 +7,7 @@ import (
 	"github.com/anacrolix/stm"
 	"github.com/anacrolix/stm/stmutil"
 
+	"github.com/anacrolix/dht/v2/containers"
 	"github.com/anacrolix/dht/v2/int160"
 	"github.com/anacrolix/dht/v2/krpc"
 	"github.com/anacrolix/dht/v2/traversal"
@@ -47,7 +48,7 @@ func newTraversal(input NewTraversalInput) stmTraversal {
 	t := stmTraversal{
 		targetInfohash:      targetInfohash,
 		triedAddrs:          stm.NewVar(stmutil.NewSet()),
-		nodesPendingContact: stm.NewVar(nodesByDistance(targetInfohash)),
+		nodesPendingContact: stm.NewVar(containers.NewAddrMaybeIdsByDistance(targetInfohash)),
 		addrBestIds:         stm.NewVar(stmutil.NewMap()),
 		pending:             stm.NewVar(0),
 	}

stm-traversal_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/anacrolix/dht/v2/int160"
+	"github.com/anacrolix/dht/v2/internal/testutil"
 	"github.com/anacrolix/dht/v2/krpc"
 )
 
@@ -19,7 +20,7 @@ func TestTraversal(t *testing.T) {
 	assert.False(t, stm.Atomically(func(tx *stm.Tx) interface{} { _, ok := traversal.popNextContact(tx); return ok }).(bool))
 	assert.False(t, stm.WouldBlock(stm.VoidOperation(traversal.waitFinished)))
 	stm.Atomically(stm.Compose(func() (ret []stm.Operation) {
-		for _, v := range sampleAddrMaybeIds[2:6] {
+		for _, v := range testutil.SampleAddrMaybeIds[2:6] {
 			ret = append(ret, traversal.pendContact(v))
 		}
 		return