@@ -101,7 +101,7 @@ data Context = Context
101 ---------------------------------------------------------------------------
101 ---------------------------------------------------------------------------
103 -- a -> AsyncT m b
103 -- a -> AsyncT m b
104− continuation :: Any
104+ continuations :: [Any]
106 -- When we suspend we save the logs in this IORef and exit.
106 -- When we suspend we save the logs in this IORef and exit.
107 , logsRef :: Maybe (IORef [Recording])
107 , logsRef :: Maybe (IORef [Recording])
@@ -161,7 +161,7 @@ initContext
161 -> Maybe (IORef [Recording])
161 -> Maybe (IORef [Recording])
162 -> Context
162 -> Context
163initContext childChan pending credit finalizer lref =
163initContext childChan pending credit finalizer lref =
164− Context { continuation = unsafeCoerce finalizer
164+ Context { continuations = [unsafeCoerce finalizer]
165 , logsRef = lref
165 , logsRef = lref
166 , location = Worker
166 , location = Worker
167 , childChannel = childChan
167 , childChannel = childChan
@@ -203,13 +203,14 @@ type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
203instance MonadAsync m => Monad (AsyncT m) where
203instance MonadAsync m => Monad (AsyncT m) where
204 return a = AsyncT . return $ Just a
204 return a = AsyncT . return $ Just a
205 m >>= f = AsyncT $ do
205 m >>= f = AsyncT $ do
206− modify $ \Context { continuation = g, .. } ->
206+ modify $ \Context {continuations = fs, ..} ->
207− let k x = f x >>= unsafeCoerce g
207+ Context {continuations = (unsafeCoerce f) : fs, ..}
208− in Context { continuation = unsafeCoerce k, .. }
210 x <- runAsyncT m
208 x <- runAsyncT m
211− Context { continuation = g } <- get
209+ modify $ \Context {continuations = (_ : fs), ..} ->
212− runAsyncT $ maybe mzero (unsafeCoerce g) x
210+ Context {continuations = fs, ..}
212+ Just y -> runAsyncT $ f y
213+ Nothing -> return Nothing
214------------------------------------------------------------------------------
215------------------------------------------------------------------------------
215-- Functor
216-- Functor
@@ -230,13 +231,37 @@ instance MonadAsync m => Applicative (AsyncT m) where
230-- Pick up from where we left in the previous thread
231-- Pick up from where we left in the previous thread
231------------------------------------------------------------------------------
232------------------------------------------------------------------------------
233−-- | Run the continuation of an action
234+-- {-# INLINE continue #-}
235+continue :: Monad m => [Any] -> a -> StateT Context m (Maybe a)
240+ modify $ \ctx -> ctx {continuations = []}
241+ runAsyncT $ (unsafeCoerce f) x
243+continue (f:fs) x = do
244+ modify $ \ctx -> ctx {continuations = fs}
245+ y <- runAsyncT $ (unsafeCoerce f) x
248+ Just z -> continue fs z
250+-- {-# INLINE runContext #-}
234runContext :: MonadAsync m => Context -> AsyncT m a -> StateT Context m ()
251runContext :: MonadAsync m => Context -> AsyncT m a -> StateT Context m ()
235runContext ctx action = do
252runContext ctx action = do
236− let s = runAsyncT $ action >>= unsafeCoerce (continuation ctx)
237− _ <- lift $ runStateT s ctx
253+ _ <- lift $ runStateT m ctx
238 return ()
254 return ()
259+ let conts = continuations ctx
260+ x <- runAsyncT action
262+ Just y -> continue conts y >> return ()
263+ Nothing -> return ()
240------------------------------------------------------------------------------
265------------------------------------------------------------------------------
241-- Thread Management (creation, reaping and killing)
266-- Thread Management (creation, reaping and killing)
242------------------------------------------------------------------------------
267------------------------------------------------------------------------------