1 files changed, 9 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 51b40b2c..12fc29f6 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -83,8 +83,8 @@ newSearch :: ( Ord addr
              Search nid addr tok ni r
              -> nid
              -> [ni]                   -- Initial nodes to query.
-             -> IO (SearchState nid addr tok ni r)
-newSearch (Search space nAddr qry) target ns = atomically $ do
+             -> STM (SearchState nid addr tok ni r)
+newSearch (Search space nAddr qry) target ns = do
     c <- newTVar 0
     q <- newTVar $ MM.fromList
                  $ map (\n -> n :-> kademliaXor space target (kademliaLocation space n))
@@ -102,7 +102,6 @@ searchK = 8
 
 sendQuery :: forall addr nid tok ni r.
             ( Ord addr
-            , Ord r
             , PSQKey nid
             , PSQKey ni
             , Show nid
@@ -144,8 +143,7 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
             [] -> return ()
 
 
-searchIsFinished :: ( Ord addr
-                    , PSQKey nid
+searchIsFinished :: ( PSQKey nid
                     , PSQKey ni
                     ) => SearchState nid addr tok ni r -> STM Bool
 searchIsFinished SearchState{ ..} = do
@@ -170,13 +168,13 @@ search ::
     , PSQKey ni
     , Show nid
     ) => Search nid addr tok ni r -> R.BucketList ni -> nid -> (r -> STM Bool) -> IO (SearchState nid addr tok ni r)
-search sch@Search{..} buckets target result = do
-    let ns = R.kclosest searchSpace searchK target buckets
-    st <- newSearch sch target ns
-    fork $ go st
+search sch buckets target result = do
+    let ns = R.kclosest (searchSpace sch) searchK target buckets
+    st <- atomically $ newSearch sch target ns
+    fork $ searchLoop sch target result st
     return st
- where
-  go s@SearchState{..} = do
+
+searchLoop sch@Search{..} target result s@SearchState{..} = do
     myThreadId >>= flip labelThread ("search."++show target)
     withTaskGroup searchAlpha $ \g -> fix $ \again -> do
         join $ atomically $ do