@@ -101,7 +101,7 @@ data Context = Context
101101 ---------------------------------------------------------------------------
102102
103103 -- a -> AsyncT m b
104− continuation :: Any
104+ continuations :: [Any]
105105
106106 -- When we suspend we save the logs in this IORef and exit.
107107 , logsRef :: Maybe (IORef [Recording])
@@ -161,7 +161,7 @@ initContext
161161 -> Maybe (IORef [Recording])
162162 -> Context
163163initContext childChan pending credit finalizer lref =
164− Context { continuation = unsafeCoerce finalizer
164+ Context { continuations = [unsafeCoerce finalizer]
165165 , logsRef = lref
166166 , location = Worker
167167 , childChannel = childChan
@@ -203,13 +203,14 @@ type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
203203instance MonadAsync m => Monad (AsyncT m) where
204204 return a = AsyncT . return $ Just a
205205 m >>= f = AsyncT $ do
206− modify $ \Context { continuation = g, .. } ->
207− let k x = f x >>= unsafeCoerce g
208− in Context { continuation = unsafeCoerce k, .. }
209−
206+ modify $ \Context {continuations = fs, ..} ->
207+ Context {continuations = (unsafeCoerce f) : fs, ..}
210208 x <- runAsyncT m
211− Context { continuation = g } <- get
212− runAsyncT $ maybe mzero (unsafeCoerce g) x
209+ modify $ \Context {continuations = (_ : fs), ..} ->
210+ Context {continuations = fs, ..}
211+ case x of
212+ Just y -> runAsyncT $ f y
213+ Nothing -> return Nothing
213214
214215------------------------------------------------------------------------------
215216-- Functor
@@ -230,13 +231,37 @@ instance MonadAsync m => Applicative (AsyncT m) where
230231-- Pick up from where we left in the previous thread
231232------------------------------------------------------------------------------
232233
233−-- | Run the continuation of an action
234+-- {-# INLINE continue #-}
235+continue :: Monad m => [Any] -> a -> StateT Context m (Maybe a)
236+continue [] x = do
237+ return (Just x)
238+
239+continue [f] x = do
240+ modify $ \ctx -> ctx {continuations = []}
241+ runAsyncT $ (unsafeCoerce f) x
242+
243+continue (f:fs) x = do
244+ modify $ \ctx -> ctx {continuations = fs}
245+ y <- runAsyncT $ (unsafeCoerce f) x
246+ case y of
247+ Nothing -> return y
248+ Just z -> continue fs z
249+
250+-- {-# INLINE runContext #-}
234251runContext :: MonadAsync m => Context -> AsyncT m a -> StateT Context m ()
235252runContext ctx action = do
236− let s = runAsyncT $ action >>= unsafeCoerce (continuation ctx)
237− _ <- lift $ runStateT s ctx
253+ _ <- lift $ runStateT m ctx
238254 return ()
239255
256+ where
257+
258+ m = do
259+ let conts = continuations ctx
260+ x <- runAsyncT action
261+ case x of
262+ Just y -> continue conts y >> return ()
263+ Nothing -> return ()
264+
240265------------------------------------------------------------------------------
241266-- Thread Management (creation, reaping and killing)
242267------------------------------------------------------------------------------