{-
  This file is copyright (c) 2009, the Snap Framework authors,
  and Patrick Thomson (for the Airship project).
  Used under the three-clause BSD license, the text of which may be
  found in the LICENSE file in the Airship root.
-}

{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}

{-
  RST is like the RWST monad, but has no Writer instance, as Writer leaks space.
  This file is almost entirely lifted from the Snap framework's implementation.
-}

module Airship.RST
       ( RST (..)
       , evalRST
       , execRST
       , mapRST
       , withRST
       , failure
       ) where

import           Control.Applicative         (Alternative (..),
                                              Applicative (..))
import           Control.Category            ((.))
import           Control.Monad               (MonadPlus (..), ap)
import           Control.Monad.Base          (MonadBase (..))
import           Control.Monad.Reader        (MonadReader (..))
import           Control.Monad.State.Class   (MonadState (..))
import           Control.Monad.Trans         (MonadIO (..), MonadTrans (..))
import           Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),
                                              MonadTransControl (..),
                                              defaultLiftBaseWith,
                                              defaultRestoreM)
import           Data.Either
import           Prelude                     (Functor (..), Monad (..), seq,
                                              ($), ($!))


newtype RST r s e m a = RST { RST r s e m a -> r -> s -> m (Either e a, s)
runRST :: r -> s -> m (Either e a, s) }


evalRST :: Monad m => RST r s e m a -> r -> s -> m (Either e a)
evalRST :: RST r s e m a -> r -> s -> m (Either e a)
evalRST RST r s e m a
m r
r s
s = do
    (Either e a
res, s
_) <- RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s
    Either e a -> m (Either e a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either e a -> m (Either e a)) -> Either e a -> m (Either e a)
forall a b. (a -> b) -> a -> b
$! Either e a
res
{-# INLINE evalRST #-}


execRST :: Monad m => RST r s e m a -> r -> s -> m s
execRST :: RST r s e m a -> r -> s -> m s
execRST RST r s e m a
m r
r s
s = do
    (Either e a
_,!s
s') <- RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s
    s -> m s
forall (m :: * -> *) a. Monad m => a -> m a
return (s -> m s) -> s -> m s
forall a b. (a -> b) -> a -> b
$! s
s'
{-# INLINE execRST #-}


withRST :: Monad m => (r' -> r) -> RST r s e m a -> RST r' s e m a
withRST :: (r' -> r) -> RST r s e m a -> RST r' s e m a
withRST r' -> r
f RST r s e m a
m = (r' -> s -> m (Either e a, s)) -> RST r' s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r' -> s -> m (Either e a, s)) -> RST r' s e m a)
-> (r' -> s -> m (Either e a, s)) -> RST r' s e m a
forall a b. (a -> b) -> a -> b
$ \r'
r' s
s -> RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m (r' -> r
f r'
r') s
s
{-# INLINE withRST #-}


instance (Monad m) => MonadReader r (RST r s e m) where
    ask :: RST r s e m r
ask = (r -> s -> m (Either e r, s)) -> RST r s e m r
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e r, s)) -> RST r s e m r)
-> (r -> s -> m (Either e r, s)) -> RST r s e m r
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> (Either e r, s) -> m (Either e r, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e r, s) -> m (Either e r, s))
-> (Either e r, s) -> m (Either e r, s)
forall a b. (a -> b) -> a -> b
$! (r -> Either e r
forall a b. b -> Either a b
Right r
r,s
s)
    local :: (r -> r) -> RST r s e m a -> RST r s e m a
local r -> r
f RST r s e m a
m = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m (r -> r
f r
r) s
s

instance (Functor m) => Functor (RST r s e m) where
    fmap :: (a -> b) -> RST r s e m a -> RST r s e m b
fmap a -> b
f RST r s e m a
m = (r -> s -> m (Either e b, s)) -> RST r s e m b
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e b, s)) -> RST r s e m b)
-> (r -> s -> m (Either e b, s)) -> RST r s e m b
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> ((Either e a, s) -> (Either e b, s))
-> m (Either e a, s) -> m (Either e b, s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Either e a
a,s
s') -> ((a -> b) -> Either e a -> Either e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Either e a
a, s
s')) (m (Either e a, s) -> m (Either e b, s))
-> m (Either e a, s) -> m (Either e b, s)
forall a b. (a -> b) -> a -> b
$ RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s

instance Monad m => Applicative (RST r s e m) where
    pure :: a -> RST r s e m a
pure = a -> RST r s e m a
forall (m :: * -> *) a. Monad m => a -> m a
return
    <*> :: RST r s e m (a -> b) -> RST r s e m a -> RST r s e m b
(<*>) = RST r s e m (a -> b) -> RST r s e m a -> RST r s e m b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap


instance MonadPlus m => Alternative (RST r s e m) where
    empty :: RST r s e m a
empty = RST r s e m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    <|> :: RST r s e m a -> RST r s e m a -> RST r s e m a
(<|>) = RST r s e m a -> RST r s e m a -> RST r s e m a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus


instance (Monad m) => MonadState s (RST r s e m) where
    get :: RST r s e m s
get   = (r -> s -> m (Either e s, s)) -> RST r s e m s
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e s, s)) -> RST r s e m s)
-> (r -> s -> m (Either e s, s)) -> RST r s e m s
forall a b. (a -> b) -> a -> b
$ \r
_ s
s -> (Either e s, s) -> m (Either e s, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e s, s) -> m (Either e s, s))
-> (Either e s, s) -> m (Either e s, s)
forall a b. (a -> b) -> a -> b
$! (s -> Either e s
forall a b. b -> Either a b
Right s
s,s
s)
    put :: s -> RST r s e m ()
put s
x = (r -> s -> m (Either e (), s)) -> RST r s e m ()
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e (), s)) -> RST r s e m ())
-> (r -> s -> m (Either e (), s)) -> RST r s e m ()
forall a b. (a -> b) -> a -> b
$ \r
_ s
_ -> (Either e (), s) -> m (Either e (), s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e (), s) -> m (Either e (), s))
-> (Either e (), s) -> m (Either e (), s)
forall a b. (a -> b) -> a -> b
$! (() -> Either e ()
forall a b. b -> Either a b
Right (),s
x)
    state :: (s -> (a, s)) -> RST r s e m a
state s -> (a, s)
act = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
s -> do
      let (a
res, !s
s') = s -> (a, s)
act s
s
      (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e a, s) -> m (Either e a, s))
-> (Either e a, s) -> m (Either e a, s)
forall a b. (a -> b) -> a -> b
$! (a -> Either e a
forall a b. b -> Either a b
Right a
res, s
s')


mapRST :: (m (Either e a, s) -> n (Either e b, s)) -> RST r s e m a -> RST r s e n b
mapRST :: (m (Either e a, s) -> n (Either e b, s))
-> RST r s e m a -> RST r s e n b
mapRST m (Either e a, s) -> n (Either e b, s)
f RST r s e m a
m = (r -> s -> n (Either e b, s)) -> RST r s e n b
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> n (Either e b, s)) -> RST r s e n b)
-> (r -> s -> n (Either e b, s)) -> RST r s e n b
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> m (Either e a, s) -> n (Either e b, s)
f (RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s)

rwsBind :: Monad m =>
           RST r s e m a
        -> (a -> RST r s e m b)
        -> RST r s e m b
rwsBind :: RST r s e m a -> (a -> RST r s e m b) -> RST r s e m b
rwsBind RST r s e m a
m a -> RST r s e m b
f = (r -> s -> m (Either e b, s)) -> RST r s e m b
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST r -> s -> m (Either e b, s)
go
  where
    go :: r -> s -> m (Either e b, s)
go r
r !s
s = do
        (Either e a
a, !s
s')  <- RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s
        case Either e a
a of
            Left e
e  -> (Either e b, s) -> m (Either e b, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e b, s) -> m (Either e b, s))
-> (Either e b, s) -> m (Either e b, s)
forall a b. (a -> b) -> a -> b
$! (e -> Either e b
forall a b. a -> Either a b
Left e
e, s
s')
            Right a
a' ->  RST r s e m b -> r -> s -> m (Either e b, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST (a -> RST r s e m b
f a
a') r
r s
s'
{-# INLINE rwsBind #-}

instance (Monad m) => Monad (RST r s e m) where
    return :: a -> RST r s e m a
return a
a = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
s -> (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e a, s) -> m (Either e a, s))
-> (Either e a, s) -> m (Either e a, s)
forall a b. (a -> b) -> a -> b
$! (a -> Either e a
forall a b. b -> Either a b
Right a
a, s
s)
    >>= :: RST r s e m a -> (a -> RST r s e m b) -> RST r s e m b
(>>=)    = RST r s e m a -> (a -> RST r s e m b) -> RST r s e m b
forall (m :: * -> *) r s e a b.
Monad m =>
RST r s e m a -> (a -> RST r s e m b) -> RST r s e m b
rwsBind
    -- fail msg = RST $ \_ _ -> fail msg

instance (MonadPlus m) => MonadPlus (RST r s e m) where
    mzero :: RST r s e m a
mzero       = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
_ -> m (Either e a, s)
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    RST r s e m a
m mplus :: RST r s e m a -> RST r s e m a -> RST r s e m a
`mplus` RST r s e m a
n = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
m r
r s
s m (Either e a, s) -> m (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` RST r s e m a -> r -> s -> m (Either e a, s)
forall r s e (m :: * -> *) a.
RST r s e m a -> r -> s -> m (Either e a, s)
runRST RST r s e m a
n r
r s
s


instance (MonadIO m) => MonadIO (RST r s e m) where
    liftIO :: IO a -> RST r s e m a
liftIO = m a -> RST r s e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> RST r s e m a) -> (IO a -> m a) -> IO a -> RST r s e m a
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO


instance MonadTrans (RST r s e) where
    lift :: m a -> RST r s e m a
lift m a
m = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
s -> do
        a
a <- m a
m
        (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e a, s) -> m (Either e a, s))
-> (Either e a, s) -> m (Either e a, s)
forall a b. (a -> b) -> a -> b
$ s
s s -> (Either e a, s) -> (Either e a, s)
`seq` (a -> Either e a
forall a b. b -> Either a b
Right a
a, s
s)


instance MonadBase b m => MonadBase b (RST r s e m) where
    liftBase :: b α -> RST r s e m α
liftBase = m α -> RST r s e m α
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m α -> RST r s e m α) -> (b α -> m α) -> b α -> RST r s e m α
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. b α -> m α
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase


instance MonadBaseControl b m => MonadBaseControl b (RST r s e m) where
     type StM (RST r s e m) a = ComposeSt (RST r s e) m a
     liftBaseWith :: (RunInBase (RST r s e m) b -> b a) -> RST r s e m a
liftBaseWith = (RunInBase (RST r s e m) b -> b a) -> RST r s e m a
forall (t :: (* -> *) -> * -> *) (b :: * -> *) (m :: * -> *) a.
(MonadTransControl t, MonadBaseControl b m) =>
(RunInBaseDefault t m b -> b a) -> t m a
defaultLiftBaseWith
     restoreM :: StM (RST r s e m) a -> RST r s e m a
restoreM = StM (RST r s e m) a -> RST r s e m a
forall (t :: (* -> *) -> * -> *) (b :: * -> *) (m :: * -> *) a.
(MonadTransControl t, MonadBaseControl b m) =>
ComposeSt t m a -> t m a
defaultRestoreM
     {-# INLINE liftBaseWith #-}
     {-# INLINE restoreM #-}

instance MonadTransControl (RST r s e) where
    type StT (RST r s e) a = (Either e a, s)
    liftWith :: (Run (RST r s e) -> m a) -> RST r s e m a
liftWith Run (RST r s e) -> m a
f = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
r s
s -> do
        a
res <- Run (RST r s e) -> m a
f (Run (RST r s e) -> m a) -> Run (RST r s e) -> m a
forall a b. (a -> b) -> a -> b
$ \(RST r -> s -> n (Either e b, s)
g) -> r -> s -> n (Either e b, s)
g r
r s
s
        (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e a, s) -> m (Either e a, s))
-> (Either e a, s) -> m (Either e a, s)
forall a b. (a -> b) -> a -> b
$! (a -> Either e a
forall a b. b -> Either a b
Right a
res, s
s)
    restoreT :: m (StT (RST r s e) a) -> RST r s e m a
restoreT m (StT (RST r s e) a)
k = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
_ -> m (Either e a, s)
m (StT (RST r s e) a)
k
    {-# INLINE liftWith #-}
    {-# INLINE restoreT #-}

failure :: Monad m => e -> RST r s e m a
failure :: e -> RST r s e m a
failure e
e = (r -> s -> m (Either e a, s)) -> RST r s e m a
forall r s e (m :: * -> *) a.
(r -> s -> m (Either e a, s)) -> RST r s e m a
RST ((r -> s -> m (Either e a, s)) -> RST r s e m a)
-> (r -> s -> m (Either e a, s)) -> RST r s e m a
forall a b. (a -> b) -> a -> b
$ \r
_ s
s -> (Either e a, s) -> m (Either e a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Either e a, s) -> m (Either e a, s))
-> (Either e a, s) -> m (Either e a, s)
forall a b. (a -> b) -> a -> b
$! (e -> Either e a
forall a b. a -> Either a b
Left e
e, s
s)