From 767d6ca633e37bfd9c9a23429a5fce44052852c9 Mon Sep 17 00:00:00 2001
From: Brian Tiger Chow <brian.holderchow@gmail.com>
Date: Sun, 21 Sep 2014 23:04:19 -0700
Subject: [PATCH] refac(bitswap, util) extract KeySet

---
 exchange/bitswap/bitswap.go            | 38 ++-------------------
 exchange/bitswap/bitswap_test.go       |  4 +--
 exchange/bitswap/strategy/interface.go | 18 ----------
 util/key_set.go                        | 46 ++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 57 deletions(-)
 create mode 100644 util/key_set.go

diff --git a/exchange/bitswap/bitswap.go b/exchange/bitswap/bitswap.go
index cf5303297..fcc558a2c 100644
--- a/exchange/bitswap/bitswap.go
+++ b/exchange/bitswap/bitswap.go
@@ -2,7 +2,6 @@ package bitswap
 
 import (
 	"errors"
-	"sync"
 
 	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
 	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
@@ -29,9 +28,7 @@ func NetMessageSession(parent context.Context, p *peer.Peer, s bsnet.NetMessageS
 		strategy:      strategy.New(),
 		routing:       directory,
 		sender:        networkAdapter,
-		wantlist: WantList{
-			data: make(map[u.Key]struct{}),
-		},
+		wantlist:      u.NewKeySet(),
 	}
 	networkAdapter.SetDelegate(bs)
 
@@ -58,38 +55,7 @@ type bitswap struct {
 	// TODO(brian): save the strategy's state to the datastore
 	strategy strategy.Strategy
 
-	wantlist WantList
-}
-
-type WantList struct {
-	lock sync.RWMutex
-	data map[u.Key]struct{}
-}
-
-func (wl *WantList) Add(k u.Key) {
-	u.DOut("Adding %v to Wantlist\n", k.Pretty())
-	wl.lock.Lock()
-	defer wl.lock.Unlock()
-
-	wl.data[k] = struct{}{}
-}
-
-func (wl *WantList) Remove(k u.Key) {
-	u.DOut("Removing %v from Wantlist\n", k.Pretty())
-	wl.lock.Lock()
-	defer wl.lock.Unlock()
-
-	delete(wl.data, k)
-}
-
-func (wl *WantList) Keys() []u.Key {
-	wl.lock.RLock()
-	defer wl.lock.RUnlock()
-	keys := make([]u.Key, 0)
-	for k, _ := range wl.data {
-		keys = append(keys, k)
-	}
-	return keys
+	wantlist u.KeySet
 }
 
 // GetBlock attempts to retrieve a particular block from peers within the
diff --git a/exchange/bitswap/bitswap_test.go b/exchange/bitswap/bitswap_test.go
index 6ec45f21c..2173fb57f 100644
--- a/exchange/bitswap/bitswap_test.go
+++ b/exchange/bitswap/bitswap_test.go
@@ -289,9 +289,7 @@ func session(net tn.Network, rs tn.RoutingServer, id peer.ID) instance {
 		strategy:      strategy.New(),
 		routing:       htc,
 		sender:        adapter,
-		wantlist: WantList{
-			data: make(map[util.Key]struct{}),
-		},
+		wantlist:      util.NewKeySet(),
 	}
 	adapter.SetDelegate(bs)
 	return instance{
diff --git a/exchange/bitswap/strategy/interface.go b/exchange/bitswap/strategy/interface.go
index 1a0e14948..48097b027 100644
--- a/exchange/bitswap/strategy/interface.go
+++ b/exchange/bitswap/strategy/interface.go
@@ -30,21 +30,3 @@ type Strategy interface {
 
 	NumBytesReceivedFrom(*peer.Peer) uint64
 }
-
-type WantList interface {
-	// Peer returns the owner of the WantList
-	Peer() *peer.Peer
-
-	// Intersection returns the keys common to both WantLists
-	Intersection(WantList) WantList
-
-	KeySet
-}
-
-// TODO(brian): potentially move this somewhere more generic. For now, it's
-// useful in BitSwap operations.
-
-type KeySet interface {
-	Contains(u.Key) bool
-	Keys() []u.Key
-}
diff --git a/util/key_set.go b/util/key_set.go
new file mode 100644
index 000000000..28d87eee5
--- /dev/null
+++ b/util/key_set.go
@@ -0,0 +1,46 @@
+package util
+
+import (
+	"sync"
+)
+
+type KeySet interface {
+	Add(Key)
+	Remove(Key)
+	Keys() []Key
+}
+
+type ks struct {
+	lock sync.RWMutex
+	data map[Key]struct{}
+}
+
+func NewKeySet() KeySet {
+	return &ks{
+		data: make(map[Key]struct{}),
+	}
+}
+
+func (wl *ks) Add(k Key) {
+	wl.lock.Lock()
+	defer wl.lock.Unlock()
+
+	wl.data[k] = struct{}{}
+}
+
+func (wl *ks) Remove(k Key) {
+	wl.lock.Lock()
+	defer wl.lock.Unlock()
+
+	delete(wl.data, k)
+}
+
+func (wl *ks) Keys() []Key {
+	wl.lock.RLock()
+	defer wl.lock.RUnlock()
+	keys := make([]Key, 0)
+	for k, _ := range wl.data {
+		keys = append(keys, k)
+	}
+	return keys
+}