on 8a63c15 help sign in
effectful/src/Effectful/Coroutine.hs
89 lines 3.0 kB history raw
1-- | Experimental support for coroutines.
2module Effectful.Coroutine
3 ( Coroutine
4 , Status(..)
5 , runCoroutine
6 , yield
7 ) where
9import Control.Concurrent.MVar
10import Control.Exception
11import Data.Function
12import qualified Control.Concurrent as C
14import Effectful.Internal.Env
15import Effectful.Internal.Has
16import Effectful.Internal.Monad
18data Coroutine o i = forall es r. Coroutine
19 { crInput :: MVar (i, Env es)
20 , crState :: MVar (State o r)
21 , crCallerEnvSize :: Int
22 }
24data Status es o i r
25 = Done r
26 | Yielded o (i -> Eff es (Status es o i r))
28runCoroutine :: Eff (Coroutine o i : es) r -> Eff es (Status es o i r)
29runCoroutine (Eff m) = impureEff $ \es -> do
30 size <- sizeEnv es
31 mvInput <- newEmptyMVar
32 mvState <- newEmptyMVar
33 mask $ \restore -> do
34 -- Create a worker thread and continue execution there.
35 tid <- C.forkIO $ do
36 let cr = Coroutine mvInput mvState size
37 er <- try $ restore . m =<< unsafeConsEnv cr es
38 tryPutMVar mvState (either Failure Success er) >>= \case
39 False -> error "unexpected"
40 True -> pure ()
41 waitForStatus restore es size tid mvInput mvState
43yield :: Coroutine o i :> es => o -> Eff es i
44yield o = impureEff $ \es -> mask $ \restore -> do
45 Coroutine{..} <- getEnv es
46 size <- sizeEnv es
47 -- Save local part of the environment as the caller will discard it.
48 localEs <- takeLastEnv (size - crCallerEnvSize) es
49 -- Pass control to the caller.
50 tryPutMVar crState (Yield o) >>= \case
51 False -> error "unexpected"
52 True -> do
53 (i, callerEs) <- restore $ takeMVar crInput
54 -- The caller resumed, reconstruct the local environment. The environment
55 -- needs to be replaced since the one we just got might be completely
56 -- different to what we had before suspending the computation, e.g. if the
57 -- computation was resumed in a different thread.
58 unsafeReplaceEnv es =<< unsafeAppendEnv callerEs localEs
59 pure i
61----------------------------------------
62-- Internal
64data State o r where
65 Failure :: SomeException -> State o r
66 Success :: r -> State o r
67 Yield :: o -> State o r
69waitForStatus
70 :: (forall a. IO a -> IO a)
71 -> Env es
72 -> Int
73 -> C.ThreadId
74 -> MVar (i, Env es)
75 -> MVar (State o r)
76 -> IO (Status es o i r)
77waitForStatus restore0 es0 size0 tid mvInput mvState = fix $ \loop -> do
78 try @SomeException (restore0 $ takeMVar mvState) >>= \case
79 Left e -> throwTo tid e >> loop
80 Right (Failure e) -> throwIO e
81 Right (Success r) -> Done r <$ unsafeTrimEnv size0 es0
82 Right (Yield o) -> Yielded o k <$ unsafeTrimEnv size0 es0
83 where
84 k i = impureEff $ \es -> mask $ \restore -> do
85 size <- sizeEnv es
86 -- Resume suspended computation with the current environment.
87 tryPutMVar mvInput (i, es) >>= \case
88 False -> error "unexpected"
89 True -> waitForStatus restore es size tid mvInput mvState