[Git][ghc/ghc][wip/romes/strict-state-tweaks] Deriving-via one-shot strict state Monad instances

Rodrigo Mesquita (@alt-romes) gitlab at gitlab.haskell.org
Tue Jul 30 09:55:42 UTC 2024



Rodrigo Mesquita pushed to branch wip/romes/strict-state-tweaks at Glasgow Haskell Compiler / GHC


Commits:
9d2ebcbc by Rodrigo Mesquita at 2024-07-30T10:55:25+01:00
Deriving-via one-shot strict state Monad instances

A small refactor to use deriving via GHC.Utils.Monad.State.Strict
Monad instances for state Monads with unboxed/strict results which all
re-implemented the one-shot trick in the instance and used unboxed
tuples:

* CmmOptM in GHC.Cmm.GenericOpt
* RegM in GHC.CmmToAsm.Reg.Linear.State
* UniqSM in GHC.Types.Unique.Supply

- - - - -


6 changed files:

- compiler/GHC/Cmm/GenericOpt.hs
- compiler/GHC/CmmToAsm/Reg/Linear/State.hs
- compiler/GHC/Types/Unique/Supply.hs
- compiler/GHC/Utils/Monad/State/Strict.hs
- testsuite/tests/count-deps/CountDepsAst.stdout
- testsuite/tests/count-deps/CountDepsParser.stdout


Changes:

=====================================
compiler/GHC/Cmm/GenericOpt.hs
=====================================
@@ -5,6 +5,7 @@
 --
 -- -----------------------------------------------------------------------------
 
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE UnboxedTuples #-}
 
@@ -26,7 +27,8 @@ import GHC.Cmm.Opt           ( cmmMachOpFold )
 import GHC.Cmm.CLabel
 import GHC.Data.FastString
 import GHC.Unit
-import Control.Monad
+import Control.Monad.Trans.Reader
+import GHC.Utils.Monad.State.Strict as Strict
 
 -- -----------------------------------------------------------------------------
 -- Generic Cmm optimiser
@@ -67,19 +69,7 @@ pattern OptMResult x y = (# x, y #)
 {-# COMPLETE OptMResult #-}
 
 newtype CmmOptM a = CmmOptM (NCGConfig -> [CLabel] -> OptMResult a)
-    deriving (Functor)
-
-instance Applicative CmmOptM where
-    pure x = CmmOptM $ \_ imports -> OptMResult x imports
-    (<*>) = ap
-
-instance Monad CmmOptM where
-  (CmmOptM f) >>= g =
-    CmmOptM $ \config imports0 ->
-                case f config imports0 of
-                  OptMResult x imports1 ->
-                    case g x of
-                      CmmOptM g' -> g' config imports1
+    deriving (Functor, Applicative, Monad) via (ReaderT NCGConfig (Strict.State [CLabel]))
 
 instance CmmMakeDynamicReferenceM CmmOptM where
     addImport = addImportCmmOpt


=====================================
compiler/GHC/CmmToAsm/Reg/Linear/State.hs
=====================================
@@ -1,4 +1,4 @@
-{-# LANGUAGE PatternSynonyms, DeriveFunctor #-}
+{-# LANGUAGE PatternSynonyms, DeriveFunctor, DerivingVia #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE UnboxedTuples #-}
 
@@ -52,31 +52,24 @@ import GHC.Types.Unique
 import GHC.Types.Unique.Supply
 import GHC.Exts (oneShot)
 
-import Control.Monad (ap)
+import GHC.Utils.Monad.State.Strict as Strict
 
-type RA_Result freeRegs a = (# RA_State freeRegs, a #)
+type RA_Result freeRegs a = (# a, RA_State freeRegs #)
 
-pattern RA_Result :: a -> b -> (# a, b #)
-pattern RA_Result a b = (# a, b #)
+pattern RA_Result :: a -> b -> (# b, a #)
+pattern RA_Result a b = (# b, a #)
 {-# COMPLETE RA_Result #-}
 
 -- | The register allocator monad type.
 newtype RegM freeRegs a
         = RegM { unReg :: RA_State freeRegs -> RA_Result freeRegs a }
-        deriving (Functor)
+        deriving (Functor, Applicative, Monad) via (Strict.State (RA_State freeRegs))
 
 -- | Smart constructor for 'RegM', as described in Note [The one-shot state
 -- monad trick] in GHC.Utils.Monad.
 mkRegM :: (RA_State freeRegs -> RA_Result freeRegs a) -> RegM freeRegs a
 mkRegM f = RegM (oneShot f)
 
-instance Applicative (RegM freeRegs) where
-      pure a  =  mkRegM $ \s -> RA_Result s a
-      (<*>) = ap
-
-instance Monad (RegM freeRegs) where
-  m >>= k   =  mkRegM $ \s -> case unReg m s of { RA_Result s a -> unReg (k a) s }
-
 -- | Get native code generator configuration
 getConfig :: RegM a NCGConfig
 getConfig = mkRegM $ \s -> RA_Result s (ra_config s)


=====================================
compiler/GHC/Types/Unique/Supply.hs
=====================================
@@ -4,6 +4,7 @@
 -}
 
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE UnboxedTuples #-}
@@ -41,6 +42,7 @@ import Control.Monad
 import Data.Word
 import GHC.Exts( Ptr(..), noDuplicate#, oneShot )
 import Foreign.Storable
+import GHC.Utils.Monad.State.Strict as Strict
 
 #include "MachDeps.h"
 
@@ -304,6 +306,8 @@ uniqFromSupply  (MkSplitUniqSupply n _ _)  = mkUniqueGrimily n
 uniqsFromSupply (MkSplitUniqSupply n _ s2) = mkUniqueGrimily n : uniqsFromSupply s2
 takeUniqFromSupply (MkSplitUniqSupply n s1 _) = (mkUniqueGrimily n, s1)
 
+{-# INLINE splitUniqSupply #-}
+
 {-
 ************************************************************************
 *                                                                      *
@@ -320,12 +324,7 @@ pattern UniqResult x y = (# x, y #)
 
 -- | A monad which just gives the ability to obtain 'Unique's
 newtype UniqSM result = USM { unUSM :: UniqSupply -> UniqResult result }
-
--- See Note [The one-shot state monad trick] for why we don't derive this.
-instance Functor UniqSM where
-  fmap f (USM m) = mkUniqSM $ \us ->
-      case m us of
-        (# r, us' #) -> UniqResult (f r) us'
+  deriving (Functor, Applicative, Monad) via (Strict.State UniqSupply)
 
 -- | Smart constructor for 'UniqSM', as described in Note [The one-shot state
 -- monad trick].
@@ -333,17 +332,6 @@ mkUniqSM :: (UniqSupply -> UniqResult a) -> UniqSM a
 mkUniqSM f = USM (oneShot f)
 {-# INLINE mkUniqSM #-}
 
-instance Monad UniqSM where
-  (>>=) = thenUs
-  (>>)  = (*>)
-
-instance Applicative UniqSM where
-    pure = returnUs
-    (USM f) <*> (USM x) = mkUniqSM $ \us0 -> case f us0 of
-                            UniqResult ff us1 -> case x us1 of
-                              UniqResult xx us2 -> UniqResult (ff xx) us2
-    (*>) = thenUs_
-
 -- TODO: try to get rid of this instance
 instance MonadFail UniqSM where
     fail = panic
@@ -356,30 +344,12 @@ initUs init_us m = case unUSM m init_us of { UniqResult r us -> (r, us) }
 initUs_ :: UniqSupply -> UniqSM a -> a
 initUs_ init_us m = case unUSM m init_us of { UniqResult r _ -> r }
 
-{-# INLINE thenUs #-}
-{-# INLINE returnUs #-}
-{-# INLINE splitUniqSupply #-}
-
--- @thenUs@ is where we split the @UniqSupply at .
-
 liftUSM :: UniqSM a -> UniqSupply -> (a, UniqSupply)
 liftUSM (USM m) us0 = case m us0 of UniqResult a us1 -> (a, us1)
 
 instance MonadFix UniqSM where
     mfix m = mkUniqSM (\us0 -> let (r,us1) = liftUSM (m r) us0 in UniqResult r us1)
 
-thenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
-thenUs (USM expr) cont
-  = mkUniqSM (\us0 -> case (expr us0) of
-                   UniqResult result us1 -> unUSM (cont result) us1)
-
-thenUs_ :: UniqSM a -> UniqSM b -> UniqSM b
-thenUs_ (USM expr) (USM cont)
-  = mkUniqSM (\us0 -> case (expr us0) of { UniqResult _ us1 -> cont us1 })
-
-returnUs :: a -> UniqSM a
-returnUs result = mkUniqSM (\us -> UniqResult result us)
-
 getUs :: UniqSM UniqSupply
 getUs = mkUniqSM (\us0 -> case splitUniqSupply us0 of (us1,us2) -> UniqResult us1 us2)
 


=====================================
compiler/GHC/Utils/Monad/State/Strict.hs
=====================================
@@ -4,7 +4,7 @@
 -- | A state monad which is strict in its state.
 module GHC.Utils.Monad.State.Strict
   ( -- * The State monad
-    State(State)
+    State(State, State' {- for deriving via purposes only -})
   , state
   , evalState
   , execState
@@ -78,8 +78,10 @@ pattern State m <- State' m
 forceState :: (# a, s #) -> (# a, s #)
 forceState (# a, !s #) = (# a, s #)
 
+-- See Note [The one-shot state monad trick] for why we don't derive this.
 instance Functor (State s) where
   fmap f m = State $ \s -> case runState' m s  of (# x, s' #) -> (# f x, s' #)
+  {-# INLINE fmap #-}
 
 instance Applicative (State s) where
   pure x  = State $ \s -> (# x, s #)
@@ -87,10 +89,20 @@ instance Applicative (State s) where
     case runState' m s  of { (# f, s' #) ->
     case runState' n s' of { (# x, s'' #) ->
                              (# f x, s'' #) }}
+  m *> n = State $ \s ->
+    case runState' m s of { (# _, s' #) ->
+    case runState' n s' of { (# x, s'' #) ->
+                             (# x, s'' #) }}
+  {-# INLINE pure #-}
+  {-# INLINE (<*>) #-}
+  {-# INLINE (*>) #-}
 
 instance Monad (State s) where
   m >>= n = State $ \s -> case runState' m s of
     (# r, !s' #) -> runState' (n r) s'
+  (>>) = (*>)
+  {-# INLINE (>>=) #-}
+  {-# INLINE (>>) #-}
 
 state :: (s -> (a, s)) -> State s a
 state f = State $ \s -> case f s of (r, s') -> (# r, s' #)


=====================================
testsuite/tests/count-deps/CountDepsAst.stdout
=====================================
@@ -215,6 +215,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


=====================================
testsuite/tests/count-deps/CountDepsParser.stdout
=====================================
@@ -234,6 +234,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



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/9d2ebcbc069b526d3af8fd17aa99b00ca5727a9a

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/9d2ebcbc069b526d3af8fd17aa99b00ca5727a9a
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20240730/94253ca2/attachment-0001.html>


More information about the ghc-commits mailing list