2 files changed, 26 insertions, 1 deletions

diff --git a/src/Network/BitTorrent/Client/Types.hs b/src/Network/BitTorrent/Client/Types.hs
index 9bae7dc3..7f228276 100644
--- a/src/Network/BitTorrent/Client/Types.hs
+++ b/src/Network/BitTorrent/Client/Types.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE TypeFamilies      #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
@@ -106,6 +107,18 @@ newtype BitTorrent a = BitTorrent
 class MonadBitTorrent m where
   liftBT :: BitTorrent a -> m a
 
+#if MIN_VERSION_monad_control(1,0,0)
+newtype BTStM a = BTStM { unBTSt :: StM (ReaderT Client IO) a }
+
+instance MonadBaseControl IO BitTorrent where
+  type StM BitTorrent a = BTStM a
+  liftBaseWith cc = BitTorrent $ liftBaseWith $ \ cc' ->
+      cc $ \ (BitTorrent m) -> BTStM <$> cc' m
+  {-# INLINE liftBaseWith #-}
+
+  restoreM = BitTorrent . restoreM . unBTSt
+  {-# INLINE restoreM #-}
+#else
 instance MonadBaseControl IO BitTorrent where
   newtype StM BitTorrent a = StM { unSt :: StM (ReaderT Client IO) a }
   liftBaseWith cc = BitTorrent $ liftBaseWith $ \ cc' ->
@@ -114,6 +127,7 @@ instance MonadBaseControl IO BitTorrent where
 
   restoreM = BitTorrent . restoreM . unSt
   {-# INLINE restoreM #-}
+#endif
 
 -- | NOP.
 instance MonadBitTorrent BitTorrent where
@@ -146,4 +160,4 @@ runBitTorrent client action = runReaderT (unBitTorrent action) client
 
 getClient :: BitTorrent Client
 getClient = BitTorrent ask
-{-# INLINE getClient #-}
\ No newline at end of file
+{-# INLINE getClient #-}
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 38b3ed11..bc5c6201 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -10,6 +10,7 @@
 --   to. Normally, you don't need to import this module, use
 --   "Network.BitTorrent.DHT" instead.
 --
+{-# LANGUAGE CPP                        #-}
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE FlexibleContexts           #-}
 {-# LANGUAGE FlexibleInstances          #-}
@@ -258,10 +259,20 @@ newtype DHT ip a = DHT { unDHT :: ReaderT (Node ip) IO a }
            , MonadBase IO, MonadReader (Node ip), MonadThrow
            )
 
+#if MIN_VERSION_monad_control(1,0,0)
+newtype DHTStM ip a = StM {
+      unSt :: StM (ReaderT (Node ip) IO) a
+    }
+#endif
+
 instance MonadBaseControl IO (DHT ip) where
+#if MIN_VERSION_monad_control(1,0,0)
+  type StM (DHT ip) a = DHTStM ip a
+#else
   newtype StM (DHT ip) a = StM {
       unSt :: StM (ReaderT (Node ip) IO) a
     }
+#endif
   liftBaseWith cc = DHT $ liftBaseWith $ \ cc' ->
       cc $ \ (DHT m) -> StM <$> cc' m
   {-# INLINE liftBaseWith #-}