From 4d4894c6f108fc4f459ea96b567254d422260c7d Mon Sep 17 00:00:00 2001
From: Harendra Kumar <harendra.kumar@gmail.com>
Date: Fri, 21 Jul 2017 14:13:55 +0530
Subject: [PATCH] Move the thread contol primitive to threads module

---
 src/Asyncly/AsyncT.hs  | 16 +++-------------
 src/Asyncly/Threads.hs | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/Asyncly/AsyncT.hs b/src/Asyncly/AsyncT.hs
index 300dd60ff..eee7e7a68 100644
--- a/src/Asyncly/AsyncT.hs
+++ b/src/Asyncly/AsyncT.hs
@@ -46,16 +46,16 @@ import           Control.Monad.Base          (MonadBase (..), liftBaseDefault)
 import           Control.Monad.Catch         (MonadThrow, throwM)
 import           Control.Monad.State         (MonadIO (..), MonadPlus (..),
                                               StateT (..), liftM, runStateT,
-                                              gets, modify)
+                                              modify)
 import           Control.Monad.Trans.Class   (MonadTrans (lift))
 import           Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),
                                               MonadTransControl (..),
                                               defaultLiftBaseWith,
                                               defaultRestoreM, liftBaseWith)
+import           Data.IORef                  (newIORef, readIORef, writeIORef)
 import           Unsafe.Coerce               (unsafeCoerce)
 
 import           Control.Monad.Trans.Recorder (MonadRecorder(..))
-import           Data.IORef                  (newIORef, readIORef, writeIORef)
 --import           Debug.Trace (traceM)
 
 import Asyncly.Threads
@@ -287,18 +287,8 @@ makeAsync = AsyncT . makeCont
 -- Controlling thread quota
 ------------------------------------------------------------------------------
 
--- XXX Should n be Word32 instead?
 -- | Runs a computation under a given thread limit.  A limit of 0 means all new
 -- tasks start synchronously in the current thread unless overridden by
 -- 'async'.
 threads :: MonadAsync m => Int -> AsyncT m a -> AsyncT m a
-threads n process = AsyncT $ do
-   oldCr <- gets threadCredit
-   newCr <- liftIO $ newIORef n
-   modify $ \s -> s { threadCredit = newCr }
-   r <- runAsyncT $ process
-        >>* (AsyncT $ do
-            modify $ \s -> s { threadCredit = oldCr }
-            return (Just ())
-            ) -- restore old credit
-   return r
+threads n action = AsyncT $ threadCtl n (runAsyncT action)
diff --git a/src/Asyncly/Threads.hs b/src/Asyncly/Threads.hs
index 7e46eedc8..c7b00fe73 100644
--- a/src/Asyncly/Threads.hs
+++ b/src/Asyncly/Threads.hs
@@ -18,6 +18,7 @@ module Asyncly.Threads
     , initContext
     , runAsyncTask
     , makeCont
+    , threadCtl
     -- , Location(..)
     -- , getLocation
     -- , setLocation
@@ -429,3 +430,23 @@ makeCont cbsetter = do
             _ <- runStateT s ctx
             return ()
     spawningParentDone
+
+------------------------------------------------------------------------------
+-- Controlling thread quota
+------------------------------------------------------------------------------
+
+-- XXX Should n be Word32 instead?
+-- | Runs a computation under a given thread limit.  A limit of 0 means all new
+-- tasks start synchronously in the current thread unless overridden by
+-- 'async'.
+threadCtl :: MonadAsync m
+    => Int
+    -> StateT Context m (Maybe a)
+    -> StateT Context m (Maybe a)
+threadCtl n action = do
+   oldCr <- gets threadCredit
+   newCr <- liftIO $ newIORef n
+   modify $ \s -> s { threadCredit = newCr }
+   r <- action
+   modify $ \s -> s { threadCredit = oldCr }
+   return r