Shake Action as a GADT

Monad.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving, GADTs, DeriveFunctor, ScopedTypeVariables, BangPatterns #-}

module Development.Shake.Internal.Core.Monad(
    RAW, Capture, runRAW,
    getRO, getRW, getsRO, getsRW, putRW, modifyRW,
    catchRAW, tryRAW, throwRAW,
    unmodifyRW, captureRAW,
    ) where

import Control.Exception.Extra
import Control.Monad.IO.Class
import Data.IORef
import Control.Applicative
import Control.Monad
import Prelude

#if __GLASGOW_HASKELL__ >= 800
import Control.Monad.Fail
#endif


data RAW ro rw a where
    Fmap :: (a -> b) -> RAW ro rw a -> RAW ro rw b
    Pure :: a -> RAW ro rw a
    Ap :: RAW ro rw (a -> b) -> RAW ro rw a -> RAW ro rw b
    Bind :: RAW ro rw a -> (a -> RAW ro rw b) -> RAW ro rw b
    LiftIO :: IO a -> RAW ro rw a
    GetRO :: RAW ro rw ro
    GetRW :: RAW ro rw rw
    PutRW :: rw -> RAW ro rw ()
    CaptureRAW :: Capture (Either SomeException a) -> RAW ro rw a
    CatchRAW :: RAW ro rw a -> (SomeException -> RAW ro rw a) -> RAW ro rw a

instance Functor (RAW ro rw) where
    fmap = Fmap

instance Applicative (RAW ro rw) where
    pure = Pure
    (<*>) = Ap

instance Monad (RAW ro rw) where
    return = Pure
    (>>) = (*>)
    (>>=) = Bind

instance MonadIO (RAW ro rw) where
    liftIO = LiftIO

instance MonadFail (RAW ro rw) where
    fail = liftIO . Prelude.fail

{-
newtype RAW ro rw a = RAW {fromRAW :: ReaderT (S ro rw) (ContT () IO) a}
    deriving (Functor, Applicative, Monad, MonadIO
#if __GLASGOW_HASKELL__ >= 800
             , MonadFail
#endif
    )
-}

type Capture a = (a -> IO ()) -> IO ()


-- | Run and then call a continuation.
runRAW :: ro -> rw -> RAW ro rw a -> Capture (Either SomeException a)
runRAW ro rw m k = do
    rw <- newIORef rw
    handler <- newIORef $ k . Left
    goRAW handler ro rw m (k . Right)
        `catch_` \e -> ($ e) =<< readIORef handler

goRAW :: forall ro rw a . IORef (SomeException -> IO ()) -> ro -> IORef rw -> RAW ro rw a -> Capture a
goRAW handler ro rw = go
    where
        go :: RAW ro rw b -> Capture b
        go x k = case x of
            Fmap f a -> go a (k . f)
            Pure a -> k a
            Ap f x -> go f $ \f -> go x (k . f)
            Bind a b -> go a $ \a -> go (b a) k
            LiftIO x -> k =<< x

            GetRO -> k ro
            GetRW -> k =<< readIORef rw
            PutRW x -> k =<< writeIORef rw x

            CatchRAW m hdl -> do
                old <- readIORef handler
                writeIORef handler $ \e -> do
                    writeIORef handler old
                    go (hdl e) k `catch_`
                        \e -> ($ e) =<< readIORef handler
                go m $ \x -> writeIORef handler old >> k x

            CaptureRAW f -> do
                old <- readIORef handler
                writeIORef handler throwIO
                f $ \x -> case x of
                    Left e -> old e
                    Right v -> do
                        writeIORef handler old
                        k v `catch_` \e -> ($ e) =<< readIORef handler
                        writeIORef handler throwIO


---------------------------------------------------------------------
-- STANDARD

getRO :: RAW ro rw ro
getRO = GetRO

getRW :: RAW ro rw rw
getRW = GetRW

getsRO :: (ro -> a) -> RAW ro rw a
getsRO f = fmap f getRO

getsRW :: (rw -> a) -> RAW ro rw a
getsRW f = fmap f getRW

-- | Strict version
putRW :: rw -> RAW ro rw ()
putRW !rw = PutRW rw

modifyRW :: (rw -> rw) -> RAW ro rw ()
modifyRW f = do x <- getRW; putRW $ f x


---------------------------------------------------------------------
-- EXCEPTIONS

catchRAW :: RAW ro rw a -> (SomeException -> RAW ro rw a) -> RAW ro rw a
catchRAW = CatchRAW


tryRAW :: RAW ro rw a -> RAW ro rw (Either SomeException a)
tryRAW m = catchRAW (fmap Right m) (return . Left)

throwRAW :: Exception e => e -> RAW ro rw a
throwRAW = liftIO . throwIO


---------------------------------------------------------------------
-- WEIRD STUFF

-- | Apply a modification, run an action, then undo the changes after.
unmodifyRW :: (rw -> (rw, rw -> rw)) -> RAW ro rw a -> RAW ro rw a
unmodifyRW f m = do
    (s2,undo) <- fmap f getRW
    putRW s2
    res <- m
    modifyRW undo
    return res


-- | Capture a continuation. The continuation should be called at most once.
--   Calling the same continuation, multiple times, in parallel, results in incorrect behaviour.
captureRAW :: Capture (Either SomeException a) -> RAW ro rw a
captureRAW f = CaptureRAW f