2 files changed, 21 insertions, 8 deletions

diff --git a/src/Network/BitTorrent/Client/Handle.hs b/src/Network/BitTorrent/Client/Handle.hs
index 6e2dd0dc..2da3b357 100644
--- a/src/Network/BitTorrent/Client/Handle.hs
+++ b/src/Network/BitTorrent/Client/Handle.hs
@@ -19,7 +19,7 @@ module Network.BitTorrent.Client.Handle
        ) where
 
 import Control.Applicative
-import Control.Concurrent
+import Control.Concurrent.Lifted as L
 import Control.Monad
 import Control.Monad.Trans
 import Data.Default
@@ -40,28 +40,28 @@ import Network.BitTorrent.Tracker  as Tracker
 
 allocHandle :: InfoHash -> BitTorrent Handle -> BitTorrent Handle
 allocHandle ih m = do
-  c @ Client {..} <- getClient
-  liftIO $ modifyMVar clientTorrents $ \ handles -> do
+  Client {..} <- getClient
+  modifyMVar clientTorrents $ \ handles -> do
     case HM.lookup ih handles of
       Just h  -> return (handles, h)
       Nothing -> do
-        h <- runBitTorrent c m
+        h <- m
         return (HM.insert ih h handles, h)
 
 freeHandle :: InfoHash -> BitTorrent () -> BitTorrent ()
 freeHandle ih finalizer = do
   c @ Client {..} <- getClient
-  liftIO $ modifyMVar_ clientTorrents $ \ handles -> do
+  modifyMVar_ clientTorrents $ \ handles -> do
     case HM.lookup ih handles of
       Nothing -> return handles
       Just _  -> do
-        runBitTorrent c finalizer
+        finalizer
         return (HM.delete ih handles)
 
 lookupHandle :: InfoHash -> BitTorrent (Maybe Handle)
 lookupHandle ih = do
   Client {..} <- getClient
-  handles     <- liftIO $ readMVar clientTorrents
+  handles     <- readMVar clientTorrents
   return (HM.lookup ih handles)
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/Client/Types.hs b/src/Network/BitTorrent/Client/Types.hs
index d6b08efe..c0d50175 100644
--- a/src/Network/BitTorrent/Client/Types.hs
+++ b/src/Network/BitTorrent/Client/Types.hs
@@ -1,4 +1,6 @@
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE TypeFamilies      #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 module Network.BitTorrent.Client.Types
        ( -- * Core types
@@ -20,8 +22,10 @@ module Network.BitTorrent.Client.Types
 import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.Chan.Split
+import Control.Monad.Base
 import Control.Monad.Logger
 import Control.Monad.Reader
+import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
 import Data.Function
 import Data.HashMap.Strict as HM
@@ -82,12 +86,21 @@ data ClientEvent
 newtype BitTorrent a = BitTorrent
   { unBitTorrent :: ReaderT Client IO a
   } deriving ( Functor, Applicative, Monad
-             , MonadIO, MonadThrow, MonadUnsafeIO
+             , MonadIO, MonadThrow, MonadUnsafeIO, MonadBase IO
              )
 
 class MonadBitTorrent m where
   liftBT :: BitTorrent a -> m a
 
+instance MonadBaseControl IO BitTorrent where
+  newtype StM BitTorrent a = StM { unSt :: StM (ReaderT Client IO) a }
+  liftBaseWith cc = BitTorrent $ liftBaseWith $ \ cc' ->
+      cc $ \ (BitTorrent m) -> StM <$> cc' m
+  {-# INLINE liftBaseWith #-}
+
+  restoreM = BitTorrent . restoreM . unSt
+  {-# INLINE restoreM #-}
+
 -- | NOP.
 instance MonadBitTorrent BitTorrent where
   liftBT = id