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

13 files changed, 99 insertions, 12 deletions

diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs b/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
index 86c3590f99..9f66793a03 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
@@ -18,7 +18,7 @@ import GHC.Cmm.BlockId
 import GHC.Cmm.Dataflow.Collections
 
 import GHC.Utils.Monad
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 import GHC.Types.Unique
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Set
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs b/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
index a93b56de95..a7ffb0555f 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
@@ -45,7 +45,7 @@ import GHC.Types.Unique.Set
 import GHC.Types.Unique.FM
 import GHC.Types.Unique
 import GHC.Builtin.Uniques
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Platform
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs b/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
index d17aebd87c..9375122567 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
@@ -33,7 +33,7 @@ import GHC.Data.Graph.Directed          (flattenSCCs)
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Platform
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 import GHC.CmmToAsm.CFG
 
 import Data.List        (nub, minimumBy)
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs b/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
index 76de71d77c..ecabbbea24 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
@@ -38,7 +38,7 @@ import GHC.Platform.Reg.Class
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Set
 import GHC.Utils.Outputable
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 
 -- | Holds interesting statistics from the register allocator.
 data RegAllocStats statics instr
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs b/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
index 4d44b43492..f514dd52ce 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
@@ -17,7 +17,7 @@ import GHC.CmmToAsm.Types
 import GHC.Types.Unique.FM
 
 import GHC.Utils.Outputable
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 
 -- | Build a map of how many times each reg was alloced, clobbered, loaded etc.
 binSpillReasons
diff --git a/compiler/GHC/CmmToAsm/Reg/Liveness.hs b/compiler/GHC/CmmToAsm/Reg/Liveness.hs
index bf53ecf421..4d70533624 100644
--- a/compiler/GHC/CmmToAsm/Reg/Liveness.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Liveness.hs
@@ -65,7 +65,7 @@ import GHC.Types.Unique.Set
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Supply
 import GHC.Data.Bag
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 
 import Data.List (mapAccumL, groupBy, partition)
 import Data.Maybe
diff --git a/compiler/GHC/Core/Opt/Exitify.hs b/compiler/GHC/Core/Opt/Exitify.hs
index 2b34992d72..7da0a68989 100644
--- a/compiler/GHC/Core/Opt/Exitify.hs
+++ b/compiler/GHC/Core/Opt/Exitify.hs
@@ -41,7 +41,7 @@ import GHC.Types.Id
 import GHC.Types.Id.Info
 import GHC.Core
 import GHC.Core.Utils
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 import GHC.Builtin.Uniques
 import GHC.Types.Var.Set
 import GHC.Types.Var.Env
diff --git a/compiler/GHC/Tc/Deriv/Functor.hs b/compiler/GHC/Tc/Deriv/Functor.hs
index 4137c0e88f..b862641e01 100644
--- a/compiler/GHC/Tc/Deriv/Functor.hs
+++ b/compiler/GHC/Tc/Deriv/Functor.hs
@@ -32,7 +32,7 @@ import GHC.Utils.Panic
 import GHC.Builtin.Names
 import GHC.Types.Name.Reader
 import GHC.Types.SrcLoc
-import GHC.Utils.Monad.State
+import GHC.Utils.Monad.State.Strict
 import GHC.Tc.Deriv.Generate
 import GHC.Tc.Utils.TcType
 import GHC.Core.TyCon
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')
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index df15959944..fae95343f6 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -721,7 +721,8 @@ Library
         GHC.Utils.Logger
         GHC.Utils.Misc
         GHC.Utils.Monad
-        GHC.Utils.Monad.State
+        GHC.Utils.Monad.State.Strict
+        GHC.Utils.Monad.State.Lazy
         GHC.Utils.Outputable
         GHC.Utils.Panic
         GHC.Utils.Panic.Plain
diff --git a/testsuite/tests/parser/should_run/CountAstDeps.stdout b/testsuite/tests/parser/should_run/CountAstDeps.stdout
index 51df841ab0..e33094795a 100644
--- a/testsuite/tests/parser/should_run/CountAstDeps.stdout
+++ b/testsuite/tests/parser/should_run/CountAstDeps.stdout
@@ -1,4 +1,4 @@
-Found 257 Language.Haskell.Syntax module dependencies
+Found 258 Language.Haskell.Syntax module dependencies
 GHC.Builtin.Names
 GHC.Builtin.PrimOps
 GHC.Builtin.Types
diff --git a/testsuite/tests/parser/should_run/CountParserDeps.stdout b/testsuite/tests/parser/should_run/CountParserDeps.stdout
index bfd270bf00..4ff132de2d 100644
--- a/testsuite/tests/parser/should_run/CountParserDeps.stdout
+++ b/testsuite/tests/parser/should_run/CountParserDeps.stdout
@@ -1,4 +1,4 @@
-Found 263 GHC.Parser module dependencies
+Found 264 GHC.Parser module dependencies
 GHC.Builtin.Names
 GHC.Builtin.PrimOps
 GHC.Builtin.Types