Split GHC.Utils.Monad.State into .Strict and .Lazy

2 files changed, 87 insertions, 1 deletions

diff --git a/compiler/GHC/Utils/Monad/State.hs b/compiler/GHC/Utils/Monad/State/Lazy.hs
index 997137525b..9d42fcb2ad 100644
--- a/compiler/GHC/Utils/Monad/State.hs
+++ b/compiler/GHC/Utils/Monad/State/Lazy.hs
@@ -2,12 +2,25 @@
 {-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE PatternSynonyms #-}
 
-module GHC.Utils.Monad.State where
+-- | A lazy state monad.
+module GHC.Utils.Monad.State.Lazy
+  ( -- * The State monda
+    State(pattern 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' { runState' :: s -> (# a, s #) }
     deriving (Functor)
 
diff --git a/compiler/GHC/Utils/Monad/State/Strict.hs b/compiler/GHC/Utils/Monad/State/Strict.hs
new file mode 100644
index 0000000000..39cd6a0773
--- /dev/null
+++ b/compiler/GHC/Utils/Monad/State/Strict.hs
@@ -0,0 +1,73 @@
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE PatternSynonyms #-}
+
+-- | A state monad which is strict in its state.
+module GHC.Utils.Monad.State.Strict
+  ( -- * The State monad
+    State(pattern State)
+  , evalState
+  , execState
+  , runState
+    -- * Operations
+  , get
+  , gets
+  , put
+  , modify
+  ) where
+
+import GHC.Prelude
+
+import GHC.Exts (oneShot)
+
+-- | A state monad which is strict in the state.
+newtype State s a = State' { runState' :: s -> (# a, s #) }
+    deriving (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 State m <- State' m
+  where
+    State m = State' (oneShot $ \s -> m 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'' #)
+
+instance Monad (State s) where
+    m >>= n  = State $ \s -> case runState' m s of
+                             (# r, !s' #) -> runState' (n r) s'
+
+get :: State s s
+get = State $ \s -> (# s, s #)
+
+gets :: (s -> a) -> State s a
+gets f = State $ \s -> (# f s, s #)
+
+put :: s -> State s ()
+put s' = State $ \_ -> (# (), s' #)
+
+modify :: (s -> s) -> State s ()
+modify f = State $ \s -> (# (), f s #)
+
+
+evalState :: State s a -> s -> 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'
+
+
+runState :: State s a -> s -> (a, s)
+runState s i = case runState' s i of
+               (# a, s' #) -> (a, s')