Use GHC's State monad consistently

GHC's internal State monad benefits from oneShot annotations on its
state, allowing for more aggressive eta expansion.

We currently don't have monad transformers with the same optimisation,
so we only change uses of the pure State monad here.

See #19657 and 19380.

Metric Decrease:
    hie002

12 files changed, 24 insertions, 12 deletions

diff --git a/compiler/GHC/CmmToAsm/Dwarf/Types.hs b/compiler/GHC/CmmToAsm/Dwarf/Types.hs
index b607d1d45e..9b9015fab9 100644
--- a/compiler/GHC/CmmToAsm/Dwarf/Types.hs
+++ b/compiler/GHC/CmmToAsm/Dwarf/Types.hs
@@ -44,7 +44,7 @@ import GHC.Utils.Misc
 import GHC.CmmToAsm.Dwarf.Constants
 
 import qualified Data.ByteString as BS
-import qualified Control.Monad.Trans.State.Strict as S
+import qualified GHC.Utils.Monad.State.Strict as S
 import Control.Monad (zipWithM, join)
 import qualified Data.Map as Map
 import Data.Word
diff --git a/compiler/GHC/Core/Opt/CallerCC.hs b/compiler/GHC/Core/Opt/CallerCC.hs
index aeffe2c034..3c47da66af 100644
--- a/compiler/GHC/Core/Opt/CallerCC.hs
+++ b/compiler/GHC/Core/Opt/CallerCC.hs
@@ -20,7 +20,7 @@ import Data.Maybe
 import qualified Text.Parsec as P
 
 import Control.Applicative
-import Control.Monad.Trans.State.Strict
+import GHC.Utils.Monad.State.Strict
 import Data.Either
 import Control.Monad
 
diff --git a/compiler/GHC/Iface/Ext/Utils.hs b/compiler/GHC/Iface/Ext/Utils.hs
index 37252c43bc..20d047d150 100644
--- a/compiler/GHC/Iface/Ext/Utils.hs
+++ b/compiler/GHC/Iface/Ext/Utils.hs
@@ -40,7 +40,7 @@ import Data.Monoid
 import Data.List                  (find)
 import Data.Traversable           ( for )
 import Data.Coerce
-import Control.Monad.Trans.State.Strict hiding (get)
+import GHC.Utils.Monad.State.Strict hiding (get)
 import Control.Monad.Trans.Reader
 import qualified Data.Tree as Tree
 
@@ -187,7 +187,7 @@ initialHTS = HTS emptyTypeMap IM.empty 0
 freshTypeIndex :: State HieTypeState TypeIndex
 freshTypeIndex = do
   index <- gets freshIndex
-  modify' $ \hts -> hts { freshIndex = index+1 }
+  modify $ \hts -> hts { freshIndex = index+1 }
   return index
 
 compressTypes
@@ -217,7 +217,7 @@ getTypeIndex t
   where
     extendHTS t ht = do
       i <- freshTypeIndex
-      modify' $ \(HTS tm tt fi) ->
+      modify $ \(HTS tm tt fi) ->
         HTS (extendTypeMap tm t i) (IM.insert i ht tt) fi
       return i
 
diff --git a/compiler/GHC/Iface/Tidy/StaticPtrTable.hs b/compiler/GHC/Iface/Tidy/StaticPtrTable.hs
index 0a61f63aed..ad7c1a3ec8 100644
--- a/compiler/GHC/Iface/Tidy/StaticPtrTable.hs
+++ b/compiler/GHC/Iface/Tidy/StaticPtrTable.hs
@@ -151,7 +151,7 @@ import GHC.Types.TyThing
 import GHC.Types.ForeignStubs
 
 import Control.Monad.Trans.Class (lift)
-import Control.Monad.Trans.State
+import Control.Monad.Trans.State.Strict
 import Data.List (intercalate)
 import Data.Maybe
 import GHC.Fingerprint
diff --git a/compiler/GHC/Stg/Debug.hs b/compiler/GHC/Stg/Debug.hs
index 17cb9fdebe..77ef7910ec 100644
--- a/compiler/GHC/Stg/Debug.hs
+++ b/compiler/GHC/Stg/Debug.hs
@@ -19,7 +19,7 @@ import GHC.Driver.Session
 
 import Control.Monad (when)
 import Control.Monad.Trans.Reader
-import Control.Monad.Trans.State
+import GHC.Utils.Monad.State.Strict
 import Control.Monad.Trans.Class
 import GHC.Types.Unique.Map
 import GHC.Types.SrcLoc
diff --git a/compiler/GHC/Stg/Pipeline.hs b/compiler/GHC/Stg/Pipeline.hs
index 1801b28692..3a00a369e8 100644
--- a/compiler/GHC/Stg/Pipeline.hs
+++ b/compiler/GHC/Stg/Pipeline.hs
@@ -32,7 +32,7 @@ import GHC.Utils.Panic
 import GHC.Utils.Logger
 import Control.Monad
 import Control.Monad.IO.Class
-import Control.Monad.Trans.State.Strict
+import GHC.Utils.Monad.State.Strict
 
 newtype StgM a = StgM { _unStgM :: StateT Char IO a }
   deriving (Functor, Applicative, Monad, MonadIO)
diff --git a/compiler/GHC/Stg/Subst.hs b/compiler/GHC/Stg/Subst.hs
index 2ff09709a9..d47107bb35 100644
--- a/compiler/GHC/Stg/Subst.hs
+++ b/compiler/GHC/Stg/Subst.hs
@@ -6,7 +6,7 @@ import GHC.Prelude
 
 import GHC.Types.Id
 import GHC.Types.Var.Env
-import Control.Monad.Trans.State.Strict
+import GHC.Utils.Monad.State.Strict
 
 import GHC.Utils.Outputable
 import GHC.Utils.Misc
diff --git a/compiler/GHC/Tc/Instance/Typeable.hs b/compiler/GHC/Tc/Instance/Typeable.hs
index c187cbf6ce..a68a92e311 100644
--- a/compiler/GHC/Tc/Instance/Typeable.hs
+++ b/compiler/GHC/Tc/Instance/Typeable.hs
@@ -45,7 +45,7 @@ import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Data.FastString ( FastString, mkFastString, fsLit )
 
-import Control.Monad.Trans.State
+import Control.Monad.Trans.State.Strict
 import Control.Monad.Trans.Class (lift)
 import Data.Maybe ( isJust )
 import Data.Word( Word64 )
diff --git a/compiler/GHC/Utils/Monad/State/Lazy.hs b/compiler/GHC/Utils/Monad/State/Lazy.hs
index 9d42fcb2ad..4f9cb034df 100644
--- a/compiler/GHC/Utils/Monad/State/Lazy.hs
+++ b/compiler/GHC/Utils/Monad/State/Lazy.hs
@@ -5,7 +5,8 @@
 -- | A lazy state monad.
 module GHC.Utils.Monad.State.Lazy
   ( -- * The State monda
-    State(pattern State)
+    State(State)
+  , state
   , evalState
   , execState
   , runState
@@ -45,6 +46,10 @@ instance Monad (State s) where
     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' #)
+
 get :: State s s
 get = State $ \s -> (# s, s #)
 
diff --git a/compiler/GHC/Utils/Monad/State/Strict.hs b/compiler/GHC/Utils/Monad/State/Strict.hs
index 39cd6a0773..27a2755067 100644
--- a/compiler/GHC/Utils/Monad/State/Strict.hs
+++ b/compiler/GHC/Utils/Monad/State/Strict.hs
@@ -5,7 +5,8 @@
 -- | A state monad which is strict in its state.
 module GHC.Utils.Monad.State.Strict
   ( -- * The State monad
-    State(pattern State)
+    State(State)
+  , state
   , evalState
   , execState
   , runState
@@ -45,6 +46,10 @@ instance Monad (State s) where
     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' #)
+
 get :: State s s
 get = State $ \s -> (# s, s #)
 
diff --git a/testsuite/tests/parser/should_run/CountAstDeps.stdout b/testsuite/tests/parser/should_run/CountAstDeps.stdout
index e33094795a..d17ccda974 100644
--- a/testsuite/tests/parser/should_run/CountAstDeps.stdout
+++ b/testsuite/tests/parser/should_run/CountAstDeps.stdout
@@ -242,6 +242,7 @@ GHC.Utils.Lexeme
 GHC.Utils.Logger
 GHC.Utils.Misc
 GHC.Utils.Monad
+GHC.Utils.Monad.State.Strict
 GHC.Utils.Outputable
 GHC.Utils.Panic
 GHC.Utils.Panic.Plain
diff --git a/testsuite/tests/parser/should_run/CountParserDeps.stdout b/testsuite/tests/parser/should_run/CountParserDeps.stdout
index 4ff132de2d..c9080fbce3 100644
--- a/testsuite/tests/parser/should_run/CountParserDeps.stdout
+++ b/testsuite/tests/parser/should_run/CountParserDeps.stdout
@@ -248,6 +248,7 @@ GHC.Utils.Lexeme
 GHC.Utils.Logger
 GHC.Utils.Misc
 GHC.Utils.Monad
+GHC.Utils.Monad.State.Strict
 GHC.Utils.Outputable
 GHC.Utils.Panic
 GHC.Utils.Panic.Plain