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

1 files changed, 18 insertions, 6 deletions

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#