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

-- | A state monad which is strict in its state.
module GHC.Utils.Monad.State.Strict
  ( -- * The State monad
    State(State, State' {- for deriving via purposes only -})
  , state
  , evalState
  , execState
  , runState
    -- * Operations
  , get
  , gets
  , put
  , modify
  ) where

import GHC.Prelude

import GHC.Exts (oneShot)

{- Note [Strict State monad]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
A State monad can be strict in many ways. Which kind of strictness do we mean?

First of, since we represent the result pair as an unboxed pair, this State
monad is strict in the sense of "Control.Monad.Trans.State.Strict": The
computations and the sequencing there-of (through 'Applicative and 'Monad'
instances) are forced strictly.

Beyond the manual unboxing of one level (which CPR could achieve similarly,
yet perhaps a bit less reliably), our 'State' is even stricter than the
transformers version:
It's also strict in the state `s` (but still lazy in the value `a`). What this
means is that whenever callers examine the state component (perhaps through
'runState'), they will find that the `s` has already been evaluated.

This additional strictness maintained in a single place, by the ubiquitous
'State' pattern synonym, by forcing the state component *after* any state action
has been run. The INVARIANT is:

> Any `s` that makes it into the unboxed pair representation is evaluated.

This invariant has another nice effect: Because the evaluatedness is quite
apparent, Nested CPR will try to unbox the state component `s` nestedly if
feasible. Detecting evaluatedness of nested components is a necessary
condition for Nested CPR to trigger; see the user's guide entry on that:
https://ghc.gitlab.haskell.org/ghc/doc/users_guide/using-optimisation.html#ghc-flag--fcpr-anal

Note that this doesn't have any effects on whether Nested CPR will unbox the `a`
component (which is still lazy by default). The user still has to use the
`return $!` idiom from the user's guide to encourage Nested CPR to unbox the `a`
result of a stateful computation.
-}

-- | A state monad which is strict in the state `s`, but lazy in the value `a`.
--
-- See Note [Strict State monad] for the particular notion of strictness and
-- implementation details.
newtype State s a = State' { forall s a. State s a -> s -> (# a, s #)
runState' :: s -> (# a, s #) }

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
-- It also implements the particular notion of strictness of this monad;
-- see Note [Strict State 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 -> (# a, s #) -> (# a, s #)
forall a s. (# a, s #) -> (# a, s #)
forceState (s -> (# a, s #)
m s
s))

-- | Forces the state component of the unboxed representation pair of 'State'.
-- See Note [Strict State monad]. This is The Place doing the forcing!
forceState :: (# a, s #) -> (# a, s #)
forceState :: forall a s. (# a, s #) -> (# a, s #)
forceState (# a
a, !s
s #) = (# a
a, s
s #)

-- See Note [The one-shot state monad trick] for why we don't derive this.
instance Functor (State s) where
  fmap :: forall a b. (a -> b) -> State s a -> State s b
fmap a -> b
f State s a
m = (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
x, s
s' #) -> (# a -> b
f a
x, s
s' #)
  {-# INLINE fmap #-}

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'' #) }}
  State s a
m *> :: forall a b. State s a -> State s b -> State s b
*> 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
_, s
s' #) ->
    case State s b -> s -> (# b, s #)
forall s a. State s a -> s -> (# a, s #)
runState' State s b
n s
s' of { (# b
x, s
s'' #) ->
                             (# b
x, s
s'' #) }}
  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}
  {-# INLINE (*>) #-}

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'
  >> :: 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 (>>=) #-}
  {-# INLINE (>>) #-}

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')