{-# LANGUAGE Strict #-}
{-# LANGUAGE UndecidableInstances #-}

module Agda.Utils.StrictReader (
    MonadReader(..)
  , module Agda.Utils.StrictReader
  , asks
  ) where

import GHC.Exts (oneShot)
import Control.Monad.Reader (MonadReader(..), asks)
import Control.Monad.State (MonadState(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.Control (MonadTransControl(..))
import Agda.Utils.ExpandCase

newtype ReaderT r m a = ReaderT {forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT :: r -> m a}

instance Monad m => Functor (ReaderT r m) where
  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> ReaderT r m a -> ReaderT r m b
fmap a -> b
f (ReaderT r -> m a
ma) = (r -> m b) -> ReaderT r m b
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m b) -> r -> m b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    a <- r -> m a
ma r
r
    pure $! f a)
  {-# INLINE (<$) #-}
  <$ :: forall a b. a -> ReaderT r m b -> ReaderT r m a
(<$) a
a (ReaderT r -> m b
ma) = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT \r
r -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

instance Monad m => Applicative (ReaderT r m) where
  {-# INLINE pure #-}
  pure :: forall a. a -> ReaderT r m a
pure a
a = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
_ -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)
  {-# INLINE (<*>) #-}
  ReaderT r -> m (a -> b)
mf <*> :: forall a b. ReaderT r m (a -> b) -> ReaderT r m a -> ReaderT r m b
<*> ReaderT r -> m a
ma = (r -> m b) -> ReaderT r m b
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m b) -> r -> m b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    f <- r -> m (a -> b)
mf r
r
    a <- ma r
    pure $! f a)
  {-# INLINE (*>) #-}
  ReaderT r -> m a
ma *> :: forall a b. ReaderT r m a -> ReaderT r m b -> ReaderT r m b
*> ReaderT r -> m b
mb = (r -> m b) -> ReaderT r m b
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m b) -> r -> m b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    _ <- r -> m a
ma r
r
    mb r)
  {-# INLINE (<*) #-}
  ReaderT r -> m a
ma <* :: forall a b. ReaderT r m a -> ReaderT r m b -> ReaderT r m a
<* ReaderT r -> m b
mb = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    a <- r -> m a
ma r
r
    _ <- mb r
    pure a)

instance Monad m => Monad (ReaderT r m) where
  {-# INLINE return #-}
  return :: forall a. a -> ReaderT r m a
return = a -> ReaderT r m a
forall a. a -> ReaderT r m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  ReaderT r -> m a
ma >>= :: forall a b. ReaderT r m a -> (a -> ReaderT r m b) -> ReaderT r m b
>>= a -> ReaderT r m b
f = (r -> m b) -> ReaderT r m b
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m b) -> r -> m b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    a <- r -> m a
ma r
r
    runReaderT (f a) r)
  {-# INLINE (>>) #-}
  >> :: forall a b. ReaderT r m a -> ReaderT r m b -> ReaderT r m b
(>>) = ReaderT r m a -> ReaderT r m b -> ReaderT r m b
forall a b. ReaderT r m a -> ReaderT r m b -> ReaderT r m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

instance Monad m => MonadReader r (ReaderT r m) where
  {-# INLINE ask #-}
  ask :: ReaderT r m r
ask = (r -> m r) -> ReaderT r m r
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m r) -> r -> m r
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> m r
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure r
r)
  {-# INLINE local #-}
  local :: forall a. (r -> r) -> ReaderT r m a -> ReaderT r m a
local r -> r
f (ReaderT r -> m a
ma) = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> m a
ma (r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
$! r -> r
f r
r)
  {-# INLINE reader #-}
  reader :: forall a. (r -> a) -> ReaderT r m a
reader r -> a
f = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$! r -> a
f r
r)

instance (Semigroup a, Monad m) => Semigroup (ReaderT r m a) where
  {-# INLINE (<>) #-}
  ReaderT r -> m a
mx <> :: ReaderT r m a -> ReaderT r m a -> ReaderT r m a
<> ReaderT r -> m a
my = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> do
    ~x <- r -> m a
mx r
r
    ~y <- my r
    pure $! x <> y)

instance (Monoid a, Monad m) => Monoid (ReaderT r m a) where
  {-# INLINE mempty #-}
  mempty :: ReaderT r m a
mempty = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty)

instance MonadTrans (ReaderT r) where
  {-# INLINE lift #-}
  lift :: forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
lift = \m a
ma -> (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT (\r
_ -> m a
ma)

instance MonadState s m => MonadState s (ReaderT r m) where
  {-# INLINE get #-}
  get :: ReaderT r m s
get = m s -> ReaderT r m s
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m s
forall s (m :: * -> *). MonadState s m => m s
get
  {-# INLINE put #-}
  put :: s -> ReaderT r m ()
put = \s
s -> m () -> ReaderT r m ()
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (s -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put s
s)

instance MonadIO m => MonadIO (ReaderT r m) where
  {-# INLINE liftIO #-}
  liftIO :: forall a. IO a -> ReaderT r m a
liftIO IO a
ma = m a -> ReaderT r m a
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO a
ma)

instance MonadTransControl (ReaderT r) where
    type StT (ReaderT r) a = a
    {-# INLINE liftWith #-}
    liftWith :: forall (m :: * -> *) a.
Monad m =>
(Run (ReaderT r) -> m a) -> ReaderT r m a
liftWith Run (ReaderT r) -> m a
f = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT (\r
r -> Run (ReaderT r) -> m a
f (\ReaderT r n b
t -> ReaderT r n b -> r -> n b
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT r n b
t r
r))
    {-# INLINE restoreT #-}
    restoreT :: forall (m :: * -> *) a.
Monad m =>
m (StT (ReaderT r) a) -> ReaderT r m a
restoreT = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> ReaderT r m a)
-> (m a -> r -> m a) -> m a -> ReaderT r m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> r -> m a
forall a b. a -> b -> a
const

instance ExpandCase (m a) => ExpandCase (ReaderT r m a) where
  type Result (ReaderT r m a) = Result (m a)
  {-# INLINE expand #-}
  expand :: ((ReaderT r m a -> Result (ReaderT r m a))
 -> Result (ReaderT r m a))
-> ReaderT r m a
expand (ReaderT r m a -> Result (ReaderT r m a)) -> Result (ReaderT r m a)
k = (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> r -> m a
forall a b. (a -> b) -> a -> b
oneShot \ ~r
r ->
    forall a. ExpandCase a => ((a -> Result a) -> Result a) -> a
expand @(m a) (((m a -> Result (m a)) -> Result (m a))
-> (m a -> Result (m a)) -> Result (m a)
forall a b. (a -> b) -> a -> b
oneShot \m a -> Result (m a)
ret ->
      let r' :: r
r' = r
r in (ReaderT r m a -> Result (ReaderT r m a)) -> Result (ReaderT r m a)
k ((ReaderT r m a -> Result (m a)) -> ReaderT r m a -> Result (m a)
forall a b. (a -> b) -> a -> b
oneShot \ReaderT r m a
act -> m a -> Result (m a)
ret (ReaderT r m a -> r -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT r m a
act r
r'))))

----------------------------------------------------------------------------------------------------

newtype Reader r a = Reader {forall r a. Reader r a -> r -> a
runReader :: r -> a}

instance Functor (Reader r) where
  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> Reader r a -> Reader r b
fmap a -> b
f (Reader r -> a
ma) = (r -> b) -> Reader r b
forall r a. (r -> a) -> Reader r a
Reader ((r -> b) -> r -> b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> let x :: a
x = r -> a
ma r
r in a -> b
f a
x)
  <$ :: forall a b. a -> Reader r b -> Reader r a
(<$) a
a (Reader r -> b
ma) = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> a
a)

instance Applicative (Reader r) where
  {-# INLINE pure #-}
  pure :: forall a. a -> Reader r a
pure a
a = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot (\r
_ -> a
a))
  {-# INLINE (<*>) #-}
  Reader r -> a -> b
mf <*> :: forall a b. Reader r (a -> b) -> Reader r a -> Reader r b
<*> Reader r -> a
ma = (r -> b) -> Reader r b
forall r a. (r -> a) -> Reader r a
Reader ((r -> b) -> r -> b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> case r -> a -> b
mf r
r of a -> b
f -> case r -> a
ma r
r of a
a -> a -> b
f a
a)
  {-# INLINE (*>) #-}
  Reader r -> a
ma *> :: forall a b. Reader r a -> Reader r b -> Reader r b
*> Reader r -> b
mb = (r -> b) -> Reader r b
forall r a. (r -> a) -> Reader r a
Reader ((r -> b) -> r -> b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> b
mb r
r)
  {-# INLINE (<*) #-}
  Reader r -> a
ma <* :: forall a b. Reader r a -> Reader r b -> Reader r a
<* Reader r -> b
mb = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> a
ma r
r)

instance Monad (Reader r) where
  {-# INLINE return #-}
  return :: forall a. a -> Reader r a
return = a -> Reader r a
forall a. a -> Reader r a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  Reader r -> a
ma >>= :: forall a b. Reader r a -> (a -> Reader r b) -> Reader r b
>>= a -> Reader r b
f = (r -> b) -> Reader r b
forall r a. (r -> a) -> Reader r a
Reader ((r -> b) -> r -> b
forall a b. (a -> b) -> a -> b
oneShot \r
r -> let x :: a
x = r -> a
ma r
r in Reader r b -> r -> b
forall r a. Reader r a -> r -> a
runReader (a -> Reader r b
f a
x) r
r)
  {-# INLINE (>>) #-}
  >> :: forall a b. Reader r a -> Reader r b -> Reader r b
(>>) = Reader r a -> Reader r b -> Reader r b
forall a b. Reader r a -> Reader r b -> Reader r b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

instance MonadReader r (Reader r) where
  {-# INLINE ask #-}
  ask :: Reader r r
ask = (r -> r) -> Reader r r
forall r a. (r -> a) -> Reader r a
Reader ((r -> r) -> r -> r
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r
r)
  {-# INLINE local #-}
  local :: forall a. (r -> r) -> Reader r a -> Reader r a
local r -> r
f (Reader r -> a
ma) = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> let x :: r
x = r -> r
f r
r in r -> a
ma r
x)
  {-# INLINE reader #-}
  reader :: forall a. (r -> a) -> Reader r a
reader r -> a
f = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> a
f r
r)

instance Semigroup a => Semigroup (Reader r a) where
  {-# INLINE (<>) #-}
  Reader r -> a
f <> :: Reader r a -> Reader r a -> Reader r a
<> Reader r -> a
g = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \r
r -> r -> a
f r
r a -> a -> a
forall a. Semigroup a => a -> a -> a
<> r -> a
g r
r)

instance Monoid a => Monoid (Reader r a) where
  {-# INLINE mempty #-}
  mempty :: Reader r a
mempty = (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader \r
_ -> a
forall a. Monoid a => a
mempty

instance ExpandCase a => ExpandCase (Reader r a) where
  type Result (Reader r a) = Result a
  {-# INLINE expand #-}
  expand :: ((Reader r a -> Result (Reader r a)) -> Result (Reader r a))
-> Reader r a
expand (Reader r a -> Result (Reader r a)) -> Result (Reader r a)
k =
    (r -> a) -> Reader r a
forall r a. (r -> a) -> Reader r a
Reader ((r -> a) -> r -> a
forall a b. (a -> b) -> a -> b
oneShot \ ~r
r ->
     forall a. ExpandCase a => ((a -> Result a) -> Result a) -> a
expand @a (((a -> Result a) -> Result a) -> (a -> Result a) -> Result a
forall a b. (a -> b) -> a -> b
oneShot \a -> Result a
deflateA -> let r' :: r
r' = r
r in (Reader r a -> Result (Reader r a)) -> Result (Reader r a)
k ((Reader r a -> Result a) -> Reader r a -> Result a
forall a b. (a -> b) -> a -> b
oneShot \Reader r a
act -> a -> Result a
deflateA (Reader r a -> r -> a
forall r a. Reader r a -> r -> a
runReader Reader r a
act r
r'))))