diff --git a/addr/sorting_test.go b/addr/sorting_test.go
deleted file mode 100644
index 7dfa013..0000000
--- a/addr/sorting_test.go
+++ /dev/null
@@ -1,32 +0,0 @@
-package addr
-
-import (
-	"sort"
-	"testing"
-)
-
-func TestAddressSorting(t *testing.T) {
-	u1 := newAddrOrFatal(t, "/ip4/152.12.23.53/udp/1234/utp")
-	u2l := newAddrOrFatal(t, "/ip4/127.0.0.1/udp/1234/utp")
-	local := newAddrOrFatal(t, "/ip4/127.0.0.1/tcp/1234")
-	norm := newAddrOrFatal(t, "/ip4/6.5.4.3/tcp/1234")
-
-	l := AddrList{local, u1, u2l, norm}
-	sort.Sort(l)
-
-	if !l[0].Equal(u2l) {
-		t.Fatal("expected utp local addr to be sorted first: ", l[0])
-	}
-
-	if !l[1].Equal(u1) {
-		t.Fatal("expected utp addr to be sorted second")
-	}
-
-	if !l[2].Equal(local) {
-		t.Fatal("expected tcp localhost addr thid")
-	}
-
-	if !l[3].Equal(norm) {
-		t.Fatal("expected normal addr last")
-	}
-}
diff --git a/pstoremem/addr_book.go b/pstoremem/addr_book.go
index 92e6f4e..0be63ac 100644
--- a/pstoremem/addr_book.go
+++ b/pstoremem/addr_book.go
@@ -7,8 +7,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/libp2p/go-libp2p-peerstore/addr"
-
 	"github.com/libp2p/go-libp2p-core/peer"
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
 	"github.com/libp2p/go-libp2p-core/record"
@@ -486,18 +484,17 @@ func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []
 	mgr.subs[p] = append(mgr.subs[p], sub)
 	mgr.mu.Unlock()
 
-	sort.Sort(addr.AddrList(initial))
+	sort.Sort(addrList(initial))
 
 	go func(buffer []ma.Multiaddr) {
 		defer close(out)
 
-		sent := make(map[string]bool, len(buffer))
-		var outch chan ma.Multiaddr
-
+		sent := make(map[string]struct{}, len(buffer))
 		for _, a := range buffer {
-			sent[string(a.Bytes())] = true
+			sent[string(a.Bytes())] = struct{}{}
 		}
 
+		var outch chan ma.Multiaddr
 		var next ma.Multiaddr
 		if len(buffer) > 0 {
 			next = buffer[0]
@@ -516,11 +513,11 @@ func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []
 					next = nil
 				}
 			case naddr := <-sub.pubch:
-				if sent[string(naddr.Bytes())] {
+				if _, ok := sent[string(naddr.Bytes())]; ok {
 					continue
 				}
+				sent[string(naddr.Bytes())] = struct{}{}
 
-				sent[string(naddr.Bytes())] = true
 				if next == nil {
 					next = naddr
 					outch = out
@@ -532,7 +529,6 @@ func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []
 				return
 			}
 		}
-
 	}(initial)
 
 	return out
diff --git a/addr/sorting.go b/pstoremem/sorting.go
similarity index 76%
rename from addr/sorting.go
rename to pstoremem/sorting.go
index 61ff4db..f36ac41 100644
--- a/addr/sorting.go
+++ b/pstoremem/sorting.go
@@ -1,4 +1,4 @@
-package addr
+package pstoremem
 
 import (
 	"bytes"
@@ -12,27 +12,20 @@ func isFDCostlyTransport(a ma.Multiaddr) bool {
 	return mafmt.TCP.Matches(a)
 }
 
-type AddrList []ma.Multiaddr
+type addrList []ma.Multiaddr
 
-func (al AddrList) Len() int {
-	return len(al)
-}
-
-func (al AddrList) Swap(i, j int) {
-	al[i], al[j] = al[j], al[i]
-}
+func (al addrList) Len() int      { return len(al) }
+func (al addrList) Swap(i, j int) { al[i], al[j] = al[j], al[i] }
 
-func (al AddrList) Less(i, j int) bool {
+func (al addrList) Less(i, j int) bool {
 	a := al[i]
 	b := al[j]
 
 	// dial localhost addresses next, they should fail immediately
 	lba := manet.IsIPLoopback(a)
 	lbb := manet.IsIPLoopback(b)
-	if lba {
-		if !lbb {
-			return true
-		}
+	if lba && !lbb {
+		return true
 	}
 
 	// dial utp and similar 'non-fd-consuming' addresses first
diff --git a/pstoremem/sorting_test.go b/pstoremem/sorting_test.go
new file mode 100644
index 0000000..8ea52d1
--- /dev/null
+++ b/pstoremem/sorting_test.go
@@ -0,0 +1,20 @@
+package pstoremem
+
+import (
+	"sort"
+	"testing"
+
+	ma "github.com/multiformats/go-multiaddr"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAddressSorting(t *testing.T) {
+	u1 := ma.StringCast("/ip4/152.12.23.53/udp/1234/utp")
+	u2l := ma.StringCast("/ip4/127.0.0.1/udp/1234/utp")
+	local := ma.StringCast("/ip4/127.0.0.1/tcp/1234")
+	norm := ma.StringCast("/ip4/6.5.4.3/tcp/1234")
+
+	l := addrList{local, u1, u2l, norm}
+	sort.Sort(l)
+	require.Equal(t, l, addrList{u2l, u1, local, norm})
+}