Generalize the type of runRW#

* Generalize the type of `runRW#` to allow arbitrary return types.

* Use `runRW#` to implement `Control.Monad.ST.Lazy.runST` (this
provides evidence that it actually works properly with the generalized
type).

* Adjust the type signature in the definition of `oneShot` to match
the one it is given in `MkId`.

Reviewers: simonmar, austin, bgamari, hvr

Reviewed By: bgamari

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D3012

GHC Trac Issues: #13178

3 files changed, 23 insertions, 13 deletions

diff --git a/compiler/basicTypes/MkId.hs b/compiler/basicTypes/MkId.hs
index 417a6c7869..65860d9045 100644
--- a/compiler/basicTypes/MkId.hs
+++ b/compiler/basicTypes/MkId.hs
@@ -1236,12 +1236,11 @@ runRWId = pcMiscPrelId runRWName ty info
 
     -- State# RealWorld
     stateRW = mkTyConApp statePrimTyCon [realWorldTy]
-    -- (# State# RealWorld, o #)
-    ret_ty  = mkTupleTy Unboxed [stateRW, openAlphaTy]
-    -- State# RealWorld -> (# State# RealWorld, o #)
+    -- o
+    ret_ty  = openAlphaTy
+    -- State# RealWorld -> o
     arg_ty  = stateRW `mkFunTy` ret_ty
-    -- (State# RealWorld -> (# State# RealWorld, o #))
-    --   -> (# State# RealWorld, o #)
+    -- (State# RealWorld -> o) -> o
     ty      = mkSpecForAllTys [runtimeRep1TyVar, openAlphaTyVar] $
               arg_ty `mkFunTy` ret_ty
 
diff --git a/libraries/base/Control/Monad/ST/Lazy/Imp.hs b/libraries/base/Control/Monad/ST/Lazy/Imp.hs
index 45d2219dce..414c06c8c3 100644
--- a/libraries/base/Control/Monad/ST/Lazy/Imp.hs
+++ b/libraries/base/Control/Monad/ST/Lazy/Imp.hs
@@ -89,12 +89,11 @@ instance Monad (ST s) where
            in
            k_a new_s
 
-{-# NOINLINE runST #-}
 -- | Return the value computed by a state transformer computation.
 -- The @forall@ ensures that the internal state used by the 'ST'
 -- computation is inaccessible to the rest of the program.
 runST :: (forall s. ST s a) -> a
-runST st = case st of ST the_st -> let (r,_) = the_st (S# realWorld#) in r
+runST (ST st) = runRW# (\s -> case st (S# s) of (r, _) -> r)
 
 -- | Allow the result of a state transformer computation to be used (lazily)
 -- inside the computation.
diff --git a/libraries/ghc-prim/GHC/Magic.hs b/libraries/ghc-prim/GHC/Magic.hs
index 96f1742dea..ecdffc5600 100644
--- a/libraries/ghc-prim/GHC/Magic.hs
+++ b/libraries/ghc-prim/GHC/Magic.hs
@@ -3,6 +3,8 @@
 {-# LANGUAGE NoImplicitPrelude #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE TypeInType #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -25,6 +27,7 @@ module GHC.Magic ( inline, noinline, lazy, oneShot, runRW# ) where
 
 import GHC.Prim
 import GHC.CString ()
+import GHC.Types (RuntimeRep, TYPE)
 
 -- | The call @inline f@ arranges that 'f' is inlined, regardless of
 -- its size. More precisely, the call @inline f@ rewrites to the
@@ -88,16 +91,25 @@ lazy x = x
 -- that would otherwise be shared are re-evaluated every time they are used. Otherwise,
 -- the use of `oneShot` is safe.
 --
--- 'oneShot' is open kinded, i.e. the type variables can refer to unlifted
--- types as well.
-oneShot :: (a -> b) -> (a -> b)
+-- 'oneShot' is representation polymorphic: the type variables may refer to lifted
+-- or unlifted types.
+oneShot :: forall (q :: RuntimeRep) (r :: RuntimeRep)
+                  (a :: TYPE q) (b :: TYPE r).
+           (a -> b) -> a -> b
 oneShot f = f
 -- Implementation note: This is wired in in MkId.lhs, so the code here is
 -- mostly there to have a place for the documentation.
 
--- | Apply a function to a 'RealWorld' token.
-runRW# :: (State# RealWorld -> (# State# RealWorld, o #))
-       -> (# State# RealWorld, o #)
+-- | Apply a function to a 'State# RealWorld' token. When manually applying
+-- a function to `realWorld#`, it is necessary to use `NOINLINE` to prevent
+-- semantically undesirable floating. `runRW#` is inlined, but only very late
+-- in compilation after all floating is complete.
+
+-- 'runRW#' is representation polymorphic: the result may have a lifted or
+-- unlifted type.
+
+runRW# :: forall (r :: RuntimeRep) (o :: TYPE r).
+          (State# RealWorld -> o) -> o
 -- See Note [runRW magic] in MkId
 #if !defined(__HADDOCK_VERSION__)
 runRW# m = m realWorld#