{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE PatternSynonyms #-}

-- | A lazy state monad.
module GHC.Utils.Monad.State.Lazy
  ( -- * The State monda
    State(State)
  , state
  , evalState
  , execState
  , runState
    -- * Operations
  , get
  , gets
  , put
  , modify
  ) where

import GHC.Prelude

import GHC.Exts (oneShot)

-- | A state monad which is lazy in the state.
newtype State s a = State' { forall s a. State s a -> s -> (# a, s #)
runState' :: s -> (# a, s #) }
    deriving ((forall a b. (a -> b) -> State s a -> State s b)
-> (forall a b. a -> State s b -> State s a) -> Functor (State s)
forall a b. a -> State s b -> State s a
forall a b. (a -> b) -> State s a -> State s b
forall s a b. a -> State s b -> State s a
forall s a b. (a -> b) -> State s a -> State s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall s a b. (a -> b) -> State s a -> State s b
fmap :: forall a b. (a -> b) -> State s a -> State s b
$c<$ :: forall s a b. a -> State s b -> State s a
<$ :: forall a b. a -> State s b -> State s a
Functor)

pattern State :: (s -> (# a, s #))
              -> State s a

-- This pattern synonym makes the monad eta-expand,
-- which as a very beneficial effect on compiler performance
-- See #18202.
-- See Note [The one-shot state monad trick] in GHC.Utils.Monad
pattern $mState :: forall {r} {s} {a}.
State s a -> ((s -> (# a, s #)) -> r) -> ((# #) -> r) -> r
$bState :: forall s a. (s -> (# a, s #)) -> State s a
State m <- State' m
  where
    State s -> (# a, s #)
m = (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 -> (# a, s #)) -> s -> (# a, s #))
-> (s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
$ \s
s -> s -> (# a, s #)
m s
s)

instance Applicative (State s) where
   pure :: forall a. a -> State s a
pure a
x   = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> State s a) -> (s -> (# a, s #)) -> State s a
forall a b. (a -> b) -> a -> b
$ \s
s -> (# a
x, s
s #)
   State s (a -> b)
m <*> :: forall a b. State s (a -> b) -> State s a -> State s b
<*> State s a
n  = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> State s b) -> (s -> (# b, s #)) -> State s b
forall a b. (a -> b) -> a -> b
$ \s
s -> case State s (a -> b) -> s -> (# a -> b, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s (a -> b)
m s
s of
                            (# a -> b
f, s
s' #) -> case State s a -> s -> (# a, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s a
n s
s' of
                                           (# a
x, s
s'' #) -> (# a -> b
f a
x, s
s'' #)

instance Monad (State s) where
    State s a
m >>= :: forall a b. State s a -> (a -> State s b) -> State s b
>>= a -> State s b
n  = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> State s b) -> (s -> (# b, s #)) -> State s b
forall a b. (a -> b) -> a -> b
$ \s
s -> case State s a -> s -> (# a, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s a
m s
s of
                             (# a
r, s
s' #) -> State s b -> s -> (# b, s #)
forall s a. State s a -> s -> (# a, s #)
runState' (a -> State s b
n a
r) s
s'

state :: (s -> (a, s)) -> State s a
state :: forall s 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 #)) -> State s a) -> (s -> (# a, s #)) -> State s a
forall a b. (a -> b) -> a -> b
$ \s
s -> case s -> (a, s)
f s
s of
                        (a
r, s
s') -> (# a
r, s
s' #)

get :: State s s
get :: forall s. State s s
get = (s -> (# s, s #)) -> State s s
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# s, s #)) -> State s s) -> (s -> (# s, s #)) -> State s s
forall a b. (a -> b) -> a -> b
$ \s
s -> (# s
s, s
s #)

gets :: (s -> a) -> State s a
gets :: forall s a. (s -> a) -> State s a
gets s -> a
f = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> State s a) -> (s -> (# a, s #)) -> State s a
forall a b. (a -> b) -> a -> b
$ \s
s -> (# s -> a
f s
s, s
s #)

put :: s -> State s ()
put :: forall s. s -> State s ()
put s
s' = (s -> (# (), s #)) -> State s ()
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# (), s #)) -> State s ())
-> (s -> (# (), s #)) -> State s ()
forall a b. (a -> b) -> a -> b
$ \s
_ -> (# (), s
s' #)

modify :: (s -> s) -> State s ()
modify :: forall s. (s -> s) -> State s ()
modify s -> s
f = (s -> (# (), s #)) -> State s ()
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# (), s #)) -> State s ())
-> (s -> (# (), s #)) -> State s ()
forall a b. (a -> b) -> a -> b
$ \s
s -> (# (), s -> s
f s
s #)


evalState :: State s a -> s -> a
evalState :: forall s a. State s a -> s -> a
evalState State s a
s s
i = case State s a -> s -> (# a, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s a
s s
i of
                (# a
a, s
_ #) -> a
a


execState :: State s a -> s -> s
execState :: forall s a. State s a -> s -> s
execState State s a
s s
i = case State s a -> s -> (# a, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s a
s s
i of
                (# a
_, s
s' #) -> s
s'


runState :: State s a -> s -> (a, s)
runState :: forall s a. State s a -> s -> (a, s)
runState State s a
s s
i = case State s a -> s -> (# a, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s a
s s
i of
               (# a
a, s
s' #) -> (a
a, s
s')