{-# LANGUAGE MagicHash #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-redundant-bang-patterns #-}

{-|
This is a plain strict state monad, where state update, monadic binding and return are all strict.
If you only need a single strict state effect, use this module.

Do not use @Control.Monad.State.Strict@ for the same purpose; it's not even strict in state updates
and is /much/ less amenable to GHC optimizations than this module.
-}

module Agda.Utils.StrictState (
    MonadState(..)
  , module Agda.Utils.StrictState
  ) where

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

newtype State s a = State {forall s a. State s a -> s -> (# a, s #)
runState# :: s -> (# a, s #)}

instance Functor (State s) where
  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> State s a -> State s b
fmap a -> b
f (State s -> (# a, s #)
g) = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
g s
s of
    (# a
a, !s
s #) -> let b :: b
b = a -> b
f a
a in (# b
b, s
s #))
  {-# INLINE (<$) #-}
  <$ :: forall a b. a -> State s b -> State s a
(<$) a
a State s b
m = (\b
_ -> a
a) (b -> a) -> State s b -> State s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State s b
m

instance Applicative (State s) where
  {-# INLINE pure #-}
  pure :: forall a. a -> State s a
pure a
a = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> (# a
a, s
s #))
  {-# INLINE (<*>) #-}
  State s -> (# a -> b, s #)
mf <*> :: forall a b. State s (a -> b) -> State s a -> State s b
<*> State s -> (# a, s #)
ma = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a -> b, s #)
mf s
s of
    (# a -> b
f, s
s #) -> case s -> (# a, s #)
ma s
s of
      (# a
a, !s
s #) -> let b :: b
b = a -> b
f a
a in (# b
b, s
s #))
  {-# INLINE (*>) #-}
  State s -> (# a, s #)
ma *> :: forall a b. State s a -> State s b -> State s b
*> State s -> (# b, s #)
mb = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
_, s
s #) -> s -> (# b, s #)
mb s
s)
  {-# INLINE (<*) #-}
  State s -> (# a, s #)
ma <* :: forall a b. State s a -> State s b -> State s a
<* State s -> (# b, s #)
mb = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
a, s
s #) -> case s -> (# b, s #)
mb s
s of
      (# b
_, s
s #) -> (# a
a, s
s #))

instance Monad (State s) where
  {-# INLINE return #-}
  return :: forall a. a -> State s a
return = a -> State s a
forall a. a -> State s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  State s -> (# a, s #)
ma >>= :: forall a b. State s a -> (a -> State s b) -> State s b
>>= a -> State s b
f = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
a, s
s #) -> State s b -> s -> (# b, s #)
forall s a. State s a -> s -> (# a, s #)
runState# (a -> State s b
f a
a) s
s)
  {-# INLINE (>>) #-}
  >> :: forall a b. State s a -> State s b -> State s b
(>>) = State s a -> State s b -> State s b
forall a b. State s a -> State s b -> State s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

{-# INLINE execState #-}
execState :: State s a -> s -> s
execState :: forall s a. State s a -> s -> s
execState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of (# a
_, s
s #) -> s
s

{-# INLINE runState #-}
runState :: State s a -> s -> (a, s)
runState :: forall s a. State s a -> s -> (a, s)
runState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of (# !a
a, !s
s #) -> (a
a, s
s)

{-# INLINE evalState #-}
evalState :: State s a -> s -> a
evalState :: forall s a. State s a -> s -> a
evalState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of (# a
a, s
_ #) -> a
a

instance MonadState s (State s) where
  {-# INLINE state #-}
  state :: forall a. (s -> (a, s)) -> State s a
state = \s -> (a, s)
f -> (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
oneShot (\s
s -> case s -> (a, s)
f s
s of (!a
a, !s
s) -> (# a
a, s
s #)))
  {-# INLINE get #-}
  get :: State s s
get = (s -> (# s, s #)) -> State s s
forall s a. (s -> (# a, s #)) -> State s a
State \s
s -> (# s
s, s
s #)
  {-# INLINE put #-}
  put :: s -> State s ()
put = \s
s -> (s -> (# (), s #)) -> State s ()
forall s a. (s -> (# a, s #)) -> State s a
State \s
_ -> (# (), s
s #)

{-# INLINE gets #-}
gets :: MonadState s m => (s -> a) -> m a
gets :: forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets s -> a
f = s -> a
f (s -> a) -> m s -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m s
forall s (m :: * -> *). MonadState s m => m s
get

{-# INLINE modify #-}
modify :: MonadState s m => (s -> s) -> m ()
modify :: forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify s -> s
f = do
  s <- m s
forall s (m :: * -> *). MonadState s m => m s
get
  let s' = s -> s
f s
s
  put s'

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

newtype StateT s m a = StateT {forall s (m :: * -> *) a. StateT s m a -> s -> m (Pair a s)
runStateT# :: s -> m (Pair a s)}

instance Functor m => Functor (StateT s m) where
  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> StateT s m a -> StateT s m b
fmap a -> b
f (StateT s -> m (Pair a s)
g) = (s -> m (Pair b s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair b s)) -> s -> m (Pair b s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> (Pair a s -> Pair b s) -> m (Pair a s) -> m (Pair b s)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(a
a :!: s
s) -> a -> b
f a
a b -> s -> Pair b s
forall a b. a -> b -> Pair a b
:!: s
s) (s -> m (Pair a s)
g s
s))
  {-# INLINE (<$) #-}
  <$ :: forall a b. a -> StateT s m b -> StateT s m a
(<$) a
a StateT s m b
m = (\b
_ -> a
a) (b -> a) -> StateT s m b -> StateT s m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT s m b
m

-- We require Monad m in order to force strictness in the Applicative sequencing.
instance Monad m => Applicative (StateT s m) where
  {-# INLINE pure #-}
  pure :: forall a. a -> StateT s m a
pure a
a = (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair a s)) -> s -> m (Pair a s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> Pair a s -> m (Pair a s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a a -> s -> Pair a s
forall a b. a -> b -> Pair a b
:!: s
s))
  {-# INLINE (<*>) #-}
  StateT s -> m (Pair (a -> b) s)
mf <*> :: forall a b. StateT s m (a -> b) -> StateT s m a -> StateT s m b
<*> StateT s -> m (Pair a s)
ma = (s -> m (Pair b s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair b s)) -> s -> m (Pair b s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> do
    f :!: s <- s -> m (Pair (a -> b) s)
mf s
s
    a :!: s <- ma s
    let b = a -> b
f a
a
    pure (b :!: s))
  {-# INLINE (*>) #-}
  StateT s -> m (Pair a s)
ma *> :: forall a b. StateT s m a -> StateT s m b -> StateT s m b
*> StateT s -> m (Pair b s)
mb = (s -> m (Pair b s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair b s)) -> s -> m (Pair b s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> do
     _ :!: s <- s -> m (Pair a s)
ma s
s
     mb s)
  {-# INLINE (<*) #-}
  StateT s -> m (Pair a s)
ma <* :: forall a b. StateT s m a -> StateT s m b -> StateT s m a
<* StateT s -> m (Pair b s)
mb = (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair a s)) -> s -> m (Pair a s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> do
    a :!: s <- s -> m (Pair a s)
ma s
s
    _ :!: s <- mb s
    pure (a :!: s))

instance Monad m => Monad (StateT s m) where
  {-# INLINE return #-}
  return :: forall a. a -> StateT s m a
return = a -> StateT s m a
forall a. a -> StateT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  StateT s -> m (Pair a s)
ma >>= :: forall a b. StateT s m a -> (a -> StateT s m b) -> StateT s m b
>>= a -> StateT s m b
f = (s -> m (Pair b s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair b s)) -> s -> m (Pair b s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> do
    a :!: s <- s -> m (Pair a s)
ma s
s
    runStateT# (f a) s)
  {-# INLINE (>>) #-}
  >> :: forall a b. StateT s m a -> StateT s m b -> StateT s m b
(>>) = StateT s m a -> StateT s m b -> StateT s m b
forall a b. StateT s m a -> StateT s m b -> StateT s m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

instance MonadTrans (StateT s) where
  {-# INLINE lift #-}
  lift :: forall (m :: * -> *) a. Monad m => m a -> StateT s m a
lift m a
ma = (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair a s)) -> s -> m (Pair a s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> do a <- m a
ma; pure (a :!: s))

instance Monad m => MonadState s (StateT s m) where
  {-# INLINE state #-}
  state :: forall a. (s -> (a, s)) -> StateT s m a
state = \s -> (a, s)
f -> (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair a s)) -> s -> m (Pair a s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (a, s)
f s
s of (!a
a, !s
s) -> Pair a s -> m (Pair a s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a a -> s -> Pair a s
forall a b. a -> b -> Pair a b
:!: s
s))
  {-# INLINE get #-}
  get :: StateT s m s
get = (s -> m (Pair s s)) -> StateT s m s
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT (\s
s -> Pair s s -> m (Pair s s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s s -> s -> Pair s s
forall a b. a -> b -> Pair a b
:!: s
s))
  {-# INLINE put #-}
  put :: s -> StateT s m ()
put s
s = (s -> m (Pair () s)) -> StateT s m ()
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT (\s
_ -> Pair () s -> m (Pair () s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() () -> s -> Pair () s
forall a b. a -> b -> Pair a b
:!: s
s))

instance MonadTransControl (StateT s) where
    type StT (StateT s) a = Pair a s
    {-# INLINE liftWith #-}
    liftWith :: forall (m :: * -> *) a.
Monad m =>
(Run (StateT s) -> m a) -> StateT s m a
liftWith Run (StateT s) -> m a
f = (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT \s
s -> do
      x <- Run (StateT s) -> m a
f \StateT s n b
t -> StateT s n b -> s -> n (Pair b s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (Pair a s)
runStateT# StateT s n b
t s
s
      pure (x :!: s)
    {-# INLINE restoreT #-}
    restoreT :: forall (m :: * -> *) a.
Monad m =>
m (StT (StateT s) a) -> StateT s m a
restoreT m (StT (StateT s) a)
msa = (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT \s
_ -> m (StT (StateT s) a)
m (Pair a s)
msa

instance MonadIO m => MonadIO (StateT s m) where
  {-# INLINE liftIO #-}
  liftIO :: forall a. IO a -> StateT s m a
liftIO IO a
ma = m a -> StateT s m a
forall (m :: * -> *) a. Monad m => m a -> StateT s 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 MonadReader r m => MonadReader r (StateT s m) where
  {-# INLINE ask #-}
  ask :: StateT s m r
ask = m r -> StateT s m r
forall (m :: * -> *) a. Monad m => m a -> StateT s m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m r
forall r (m :: * -> *). MonadReader r m => m r
ask
  {-# INLINE local #-}
  local :: forall a. (r -> r) -> StateT s m a -> StateT s m a
local = \r -> r
f (StateT s -> m (Pair a s)
ma) -> (s -> m (Pair a s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (Pair a s)) -> StateT s m a
StateT ((s -> m (Pair a s)) -> s -> m (Pair a s)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> (r -> r) -> m (Pair a s) -> m (Pair a s)
forall a. (r -> r) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f (s -> m (Pair a s)
ma s
s))

{-# INLINE execStateT #-}
execStateT :: Monad m => StateT s m a -> s -> m s
execStateT :: forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (StateT s -> m (Pair a s)
f) s
s = do _ :!: s <- s -> m (Pair a s)
f s
s; pure s

{-# INLINE runStateT #-}
runStateT :: Monad m => StateT s m a -> s -> m (a, s)
runStateT :: forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m (a, s)
runStateT (StateT s -> m (Pair a s)
f) s
s = do a :!: s <- s -> m (Pair a s)
f s
s; pure (a, s)

{-# INLINE evalStateT #-}
evalStateT :: Monad m => StateT s m a -> s -> m a
evalStateT :: forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (StateT s -> m (Pair a s)
f) s
s = do a :!: _ <- s -> m (Pair a s)
f s
s; pure a

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