diff --git a/compiler/GHC/Utils/Monad/State/Strict.hs b/compiler/GHC/Utils/Monad/State/Strict.hs index d5298955a555db27d0c5fd9073e539354145a66c..b7cafb7f522ff1b89169f8c44acb0a482d75eb48 100644 --- a/compiler/GHC/Utils/Monad/State/Strict.hs +++ b/compiler/GHC/Utils/Monad/State/Strict.hs @@ -21,9 +21,45 @@ import GHC.Prelude import GHC.Exts (oneShot) --- | A state monad which is strict in the state. +{- 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' { runState' :: s -> (# a, s #) } - deriving (Functor) pattern State :: (s -> (# a, s #)) -> State s a @@ -32,47 +68,51 @@ pattern State :: (s -> (# a, s #)) -- 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 State m <- State' m where - State m = State' (oneShot $ \s -> m s) + State m = State' (oneShot $ \s -> forceState (m 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 (# a, !s #) = (# a, s #) + +instance Functor (State s) where + fmap f m = State $ \s -> case runState' m s of (# x, s' #) -> (# f x, s' #) instance Applicative (State s) where - pure x = State $ \(!s) -> (# x, s #) - m <*> n = State $ \s -> case runState' m s of - (# f, !s' #) -> case runState' n s' of - (# x, !s'' #) -> (# f x, s'' #) + pure x = State $ \s -> (# x, s #) + m <*> n = State $ \s -> + case runState' m s of { (# f, s' #) -> + case runState' n s' of { (# x, s'' #) -> + (# f x, s'' #) }} instance Monad (State s) where - m >>= n = State $ \s -> case runState' m s of - (# r, !s' #) -> runState' (n r) s' + m >>= n = State $ \s -> case runState' m s of + (# r, !s' #) -> runState' (n r) s' state :: (s -> (a, s)) -> State s a -state f = State $ \s -> case f s of - (r, !s') -> (# r, s' #) +state f = State $ \s -> case f s of (r, s') -> (# r, s' #) get :: State s s -get = State $ \(!s) -> (# s, s #) +get = State $ \s -> (# s, s #) gets :: (s -> a) -> State s a -gets f = State $ \(!s) -> (# f s, s #) +gets f = State $ \s -> (# f s, s #) put :: s -> State s () -put !s' = State $ \_ -> (# (), s' #) +put s' = State $ \_ -> (# (), s' #) modify :: (s -> s) -> State s () -modify f = State $ \s -> let !s' = f s in (# (), s' #) - +modify f = State $ \s -> (# (), f s #) evalState :: State s a -> s -> a -evalState s i = case runState' s i of - (# a, _ #) -> a - +evalState s i = case runState' s i of (# a, _ #) -> a execState :: State s a -> s -> s -execState s i = case runState' s i of - (# _, s' #) -> s' - +execState s i = case runState' s i of (# _, s' #) -> s' runState :: State s a -> s -> (a, s) -runState s i = case runState' s i of - (# a, !s' #) -> (a, s') +runState s i = case runState' s i of (# a, !s' #) -> (a, s')