1{-# LANGUAGE LambdaCase #-}
2{-# LANGUAGE GADTs, ScopedTypeVariables, TupleSections, GeneralizedNewtypeDeriving #-}
3
4module Development.Shake.Internal.Core.Monad(
5    RAW, Capture, runRAW,
6    getRO, getRW, putRW, modifyRW,
7    stepRAW,
8    catchRAW, tryRAW, throwRAW, finallyRAW,
9    captureRAW,
10    ) where
11
12import Control.Exception.Extra
13import Development.Shake.Internal.Errors
14import Control.Monad.IO.Class
15import Data.IORef
16import Control.Monad
17import System.IO
18import Data.Semigroup
19import Control.Monad.Fail
20import Prelude
21
22
23data RAW k v ro rw a where
24    Fmap :: (a -> b) -> RAW k v ro rw a -> RAW k v ro rw b
25    Pure :: a -> RAW k v ro rw a
26    Ap :: RAW k v ro rw (a -> b) -> RAW k v ro rw a -> RAW k v ro rw b
27    Next :: RAW k v ro rw a -> RAW k v ro rw b -> RAW k v ro rw b
28    Bind :: RAW k v ro rw a -> (a -> RAW k v ro rw b) -> RAW k v ro rw b
29    LiftIO :: IO a -> RAW k v ro rw a
30    GetRO :: RAW k v ro rw ro
31    GetRW :: RAW k v ro rw rw
32    PutRW :: !rw -> RAW k v ro rw ()
33    ModifyRW :: (rw -> rw) -> RAW k v ro rw ()
34    StepRAW :: k -> RAW k v ro rw v
35    CaptureRAW :: Capture (Either SomeException a) -> RAW k v ro rw a
36    CatchRAW :: RAW k v ro rw a -> (SomeException -> RAW k v ro rw a) -> RAW k v ro rw a
37
38instance Functor (RAW k v ro rw) where
39    fmap = Fmap
40
41instance Applicative (RAW k v ro rw) where
42    pure = Pure
43    (*>) = Next
44    (<*>) = Ap
45
46instance Monad (RAW k v ro rw) where
47    return = pure
48    (>>) = (*>)
49    (>>=) = Bind
50
51instance MonadIO (RAW k v ro rw) where
52    liftIO = LiftIO
53
54instance MonadFail (RAW k v ro rw) where
55    fail = liftIO . Control.Monad.Fail.fail
56
57instance Semigroup a => Semigroup (RAW k v ro rw a) where
58    (<>) a b = (<>) <$> a <*> b
59
60instance (Semigroup a, Monoid a) => Monoid (RAW k v ro rw a) where
61    mempty = pure mempty
62    mappend = (<>)
63
64
65type Capture a = (a -> IO ()) -> IO ()
66
67
68-- Useful for checking that all continuations are run only once
69-- Cannot be enabled for performance reasons and because some of
70-- "monad test" deliberately breaks the invariant to check it doesn't go wrong
71assertOnceCheck = False
72
73assertOnce :: MonadIO m => String -> (a -> m b) -> IO (a -> m b)
74assertOnce msg k
75    | not assertOnceCheck = pure k
76    | otherwise = do
77        ref <- liftIO $ newIORef False
78        pure $ \v -> do
79            liftIO $ join $ atomicModifyIORef ref $ \old -> (True,) $ when old $ do
80                hPutStrLn stderr "FATAL ERROR: assertOnce failed"
81                Prelude.fail $ "assertOnce failed: " ++ msg
82            k v
83
84-- | Run and then call a continuation.
85runRAW :: ([k] -> RAW k v ro rw [v]) -> ro -> rw -> RAW k v ro rw a -> Capture (Either SomeException a)
86runRAW step ro rw m k = do
87    k <- assertOnce "runRAW" k
88    rw <- newIORef rw
89    handler <- newIORef throwIO
90    steps <- newSteps
91    writeIORef handler $ \e -> do
92        -- make sure we never call the error continuation twice
93        writeIORef handler throwIO
94        k $ Left e
95    -- If the continuation itself throws an error we need to make sure we
96    -- don't end up running it twice (once with its result, once with its own exception)
97    goRAW step steps handler ro rw m (\v -> do writeIORef handler throwIO; k $ Right v)
98        `catch_` \e -> ($ e) =<< readIORef handler
99
100
101goRAW :: forall k v ro rw a . ([k] -> RAW k v ro rw [v]) -> Steps k v -> IORef (SomeException -> IO ()) -> ro -> IORef rw -> RAW k v ro rw a -> Capture a
102goRAW step steps handler ro rw = \x k -> go x $ \v -> sio v k
103    where
104        sio :: SIO b -> Capture b
105        sio (SIO v) k = flush $ do v <- v; k v
106
107        flush :: IO () -> IO ()
108        flush k = do
109            v <- flushSteps steps
110            case v of
111                Nothing -> k
112                Just f -> go (f step) $ const k
113
114        unflush :: IO ()
115        unflush = unflushSteps steps
116
117        go :: RAW k v ro rw b -> Capture (SIO b)
118        go x k = case x of
119            Fmap f a -> go a $ \v -> k $ fmap f v
120            Pure a -> k $ pure a
121            Ap f x -> go f $ \f -> go x $ \v -> k $ f <*> v
122            Next a b -> go a $ \a -> go b $ \b -> k $ a *> b
123            StepRAW q -> do
124                v <- addStep steps q
125                k v
126
127            Bind a b -> go a $ \a -> sio a $ \a -> go (b a) k
128            LiftIO act -> flush $ do v <- act; k $ pure v
129
130            GetRO -> k $ pure ro
131            GetRW -> flush $ k . pure =<< readIORef rw
132            PutRW x -> flush $ writeIORef rw x >> k (pure ())
133            ModifyRW f -> flush $ modifyIORef' rw f >> k (pure ())
134
135            CatchRAW m hdl -> flush $ do
136                hdl <- assertOnce "CatchRAW" hdl
137                old <- readIORef handler
138                writeIORef handler $ \e -> do
139                    writeIORef handler old
140                    go (hdl e) k `catch_`
141                        \e -> do unflush; ($ e) =<< readIORef handler
142                go m $ \x -> writeIORef handler old >> k x
143
144            CaptureRAW f -> flush $ do
145                f <- assertOnce "CaptureRAW" f
146                old <- readIORef handler
147                writeIORef handler throwIO
148                f $ \case
149                    Left e -> old e
150                    Right v -> do
151                        writeIORef handler old
152                        k (pure v) `catch_` \e -> do unflush; ($ e) =<< readIORef handler
153
154
155newtype SIO a = SIO (IO a)
156    deriving (Functor, Monad, Applicative)
157
158
159newtype Steps k v = Steps (IORef [(k, IORef v)])
160
161newSteps :: IO (Steps k v)
162newSteps = Steps <$> newIORef []
163
164addStep :: Steps k v -> k -> IO (SIO v)
165addStep (Steps ref) k = do
166    out <- newIORef $ throwImpure $ errorInternal "Monad, addStep not flushed"
167    modifyIORef ref ((k,out):)
168    pure $ SIO $ readIORef out
169
170unflushSteps :: Steps k v -> IO ()
171unflushSteps (Steps ref) = writeIORef ref []
172
173flushSteps :: MonadIO m => Steps k v -> IO (Maybe (([k] -> m [v]) -> m ()))
174flushSteps (Steps ref) = do
175    v <- reverse <$> readIORef ref
176    case v of
177        [] -> pure Nothing
178        xs -> do
179            writeIORef ref []
180            pure $ Just $ \step -> do
181                vs <- step $ map fst xs
182                liftIO $ zipWithM_ writeIORef (map snd xs) vs
183
184
185---------------------------------------------------------------------
186-- STANDARD
187
188getRO :: RAW k v ro rw ro
189getRO = GetRO
190
191getRW :: RAW k v ro rw rw
192getRW = GetRW
193
194-- | Strict version
195putRW :: rw -> RAW k v ro rw ()
196putRW = PutRW
197
198modifyRW :: (rw -> rw) -> RAW k v ro rw ()
199modifyRW = ModifyRW
200
201
202---------------------------------------------------------------------
203-- EXCEPTIONS
204
205catchRAW :: RAW k v ro rw a -> (SomeException -> RAW k v ro rw a) -> RAW k v ro rw a
206catchRAW = CatchRAW
207
208tryRAW :: RAW k v ro rw a -> RAW k v ro rw (Either SomeException a)
209tryRAW m = catchRAW (fmap Right m) (pure . Left)
210
211throwRAW :: Exception e => e -> RAW k v ro rw a
212-- Note that while we could directly pass this to the handler
213-- that would avoid triggering the catch, which would mean they built up on the stack
214throwRAW = liftIO . throwIO
215
216finallyRAW :: RAW k v ro rw a -> RAW k v ro rw b -> RAW k v ro rw a
217finallyRAW a undo = do
218    r <- catchRAW a (\e -> undo >> throwRAW e)
219    undo
220    pure r
221
222
223---------------------------------------------------------------------
224-- CONTINUATIONS
225
226-- | Capture a continuation. The continuation should be called at most once.
227--   Calling the same continuation, multiple times, in parallel, results in incorrect behaviour.
228captureRAW :: Capture (Either SomeException a) -> RAW k v ro rw a
229captureRAW = CaptureRAW
230
231
232---------------------------------------------------------------------
233-- STEPS
234
235stepRAW :: k -> RAW k v ro rw v
236stepRAW = StepRAW
237