3 files changed, 22 insertions, 82 deletions

diff --git a/src/Data/Torrent.hs b/src/Data/Torrent.hs
index a9efbfce..ae40bef4 100644
--- a/src/Data/Torrent.hs
+++ b/src/Data/Torrent.hs
@@ -71,7 +71,6 @@ import Data.Char
 import Data.Foldable
 import           Data.Map (Map)
 import qualified Data.Map as M
-import           Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import           Data.ByteString.Internal
 import qualified Data.ByteString.Char8 as BC (pack, unpack)
@@ -88,8 +87,6 @@ import Network.URI
 import System.FilePath
 import Numeric
 
-import Debug.Trace
-
 {-----------------------------------------------------------------------
     Info hash
 -----------------------------------------------------------------------}
diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 7c1e02e2..8c3189d3 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -7,88 +7,36 @@
 --
 {-# LANGUAGE RecordWildCards #-}
 module Network.BitTorrent
-       (
-         module Data.Torrent
+       ( module Data.Torrent
 
-         -- * Session
-       , ThreadCount
-       , defaultThreadCount
+       , TorrentLoc(..), Progress(..)
+       , ThreadCount, SessionCount
 
-         -- ** Client
        , ClientSession( clientPeerId, allowedExtensions )
+       , withDefaultClient, defaultThreadCount
+       , addTorrent
+       , removeTorrent
 
-       , withDefaultClient
-
-       , Progress(..)
        , getCurrentProgress
        , getPeerCount
        , getSwarmCount
-
-       , TorrentLoc(..)
-       , addTorrent
-       , removeTorrent
-
-         -- ** Swarm
-       , SwarmSession(torrentMeta)
-
-       , newLeecher
-       , newSeeder
-
-       , SessionCount
        , getSessionCount
 
-         -- * Discovery
-       , discover
-       , exchange
-
-
-         -- * Peer to Peer
-       , P2P
-
-         -- ** Session
-       , PeerSession( PeerSession, connectedPeerAddr
-                    , swarmSession, enabledExtensions
-                    )
-
-       , getHaveCount
-       , getWantCount
-       , getPieceCount
-
-
-         -- ** Transfer
-       , Block(..), ppBlock
-       , BlockIx(..), ppBlockIx
-
-         -- ** Control
-       , SessionException
-       , disconnect
-       , protocolError
-
-         -- ** Events
-       , Event(..)
-       , awaitEvent, yieldEvent
-
          -- * Extensions
-       , Extension, defaultExtensions, ppExtension
+       , Extension
+       , defaultExtensions
+       , ppExtension
        ) where
 
 import Control.Concurrent
-import Control.Concurrent.STM
-import Control.Exception
 import Control.Monad
-import Control.Monad.Reader
-
+import Control.Monad.Trans
 import Network
-
-import Data.Bitfield as BF
 import Data.Torrent
 import Network.BitTorrent.Sessions.Types
 import Network.BitTorrent.Sessions
-import Network.BitTorrent.Peer
 import Network.BitTorrent.Extension
 import Network.BitTorrent.Exchange
-import Network.BitTorrent.Exchange.Protocol
-import Network.BitTorrent.Tracker
 
 import System.Torrent.Storage
 
diff --git a/src/Network/BitTorrent/Sessions.hs b/src/Network/BitTorrent/Sessions.hs
index da57ec3b..2e51fde6 100644
--- a/src/Network/BitTorrent/Sessions.hs
+++ b/src/Network/BitTorrent/Sessions.hs
@@ -43,16 +43,6 @@ module Network.BitTorrent.Sessions
        , newSeeder
        , getClientBitfield
 
-         -- * Peer
-       , PeerSession( PeerSession, connectedPeerAddr
-                    , swarmSession, enabledExtensions
-                    , sessionState
-                    )
-       , SessionState
-       , initiatePeerSession
-       , acceptPeerSession
-       , listener
-
          -- * Timeouts
        , updateIncoming, updateOutcoming
        , discover
@@ -64,13 +54,11 @@ import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Concurrent.MSem as MSem
-import Control.Lens
 import Control.Monad (when, forever, (>=>))
 import Control.Exception
 import Control.Monad.Trans
 
 import Data.IORef
-import Data.Foldable (mapM_)
 import Data.Map as M
 import Data.HashMap.Strict as HM
 import Data.Set as S
@@ -265,10 +253,10 @@ unregisterSwarmSession SwarmSession {..} =
 
 getSwarm :: ClientSession -> InfoHash -> IO SwarmSession
 getSwarm cs @ ClientSession {..} ih = do
-  ss <- readTVarIO $ swarmSessions
+  ss <- readTVarIO swarmSessions
   case M.lookup ih ss of
     Just sw -> return sw
-    Nothing -> undefined -- openSwarm cs
+    Nothing -> undefined -- openSwarmSession cs
 
 newSeeder :: ClientSession -> Torrent -> IO SwarmSession
 newSeeder cs t @ Torrent {..}
@@ -361,6 +349,14 @@ TODO: utilize peer Id.
 TODO: use STM semaphore
 -----------------------------------------------------------------------}
 
+registerPeerSession :: PeerSession -> IO ()
+registerPeerSession ps @ PeerSession {..}  =
+  atomically $ modifyTVar' (connectedPeers swarmSession) (S.insert ps)
+
+unregisterPeerSession :: PeerSession -> IO ()
+unregisterPeerSession ps @ PeerSession {..} =
+  atomically $ modifyTVar' (connectedPeers swarmSession) (S.delete ps)
+
 openSession :: SwarmSession -> PeerAddr -> Handshake -> IO PeerSession
 openSession ss @ SwarmSession {..} addr Handshake {..} = do
   let clientCaps = encodeExts $ allowedExtensions $ clientSession
@@ -372,12 +368,11 @@ openSession ss @ SwarmSession {..} addr Handshake {..} = do
     <*> (newIORef . initialSessionState . totalCount =<< readTVarIO clientBitfield)
     -- TODO we could implement more interesting throtling scheme
     -- using connected peer information
-  atomically $ modifyTVar' connectedPeers (S.insert ps)
+  registerPeerSession ps
   return ps
 
 closeSession :: PeerSession -> IO ()
-closeSession ps @ PeerSession {..} = do
-  atomically $ modifyTVar' (connectedPeers swarmSession) (S.delete ps)
+closeSession = unregisterPeerSession
 
 type PeerConn = (Socket, PeerSession)
 type Exchange = PeerConn -> IO ()