Make oneShot open-kinded

akio wants to use oneShot with unlifted types as well, and there is no
good reason not to let him. This changes the type of the built-in
oneShot definition to open kinds, and also expand the documentation a
little bit.

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

4 files changed, 27 insertions, 2 deletions

diff --git a/compiler/basicTypes/MkId.hs b/compiler/basicTypes/MkId.hs
index 11f8f788cd..29e0e64c9b 100644
--- a/compiler/basicTypes/MkId.hs
+++ b/compiler/basicTypes/MkId.hs
@@ -1142,11 +1142,11 @@ oneShotId = pcMiscPrelId oneShotName ty info
   where
     info = noCafIdInfo `setInlinePragInfo` alwaysInlinePragma
                        `setUnfoldingInfo`  mkCompulsoryUnfolding rhs
-    ty  = mkForAllTys [alphaTyVar, betaTyVar] (mkFunTy fun_ty fun_ty)
+    ty  = mkForAllTys [openAlphaTyVar, openBetaTyVar] (mkFunTy fun_ty fun_ty)
     fun_ty = mkFunTy alphaTy betaTy
     [body, x] = mkTemplateLocals [fun_ty, alphaTy]
     x' = setOneShotLambda x
-    rhs = mkLams [alphaTyVar, betaTyVar, body, x'] $ Var body `App` Var x
+    rhs = mkLams [openAlphaTyVar, openBetaTyVar, body, x'] $ Var body `App` Var x
 
 
 --------------------------------------------------------------------------------
diff --git a/libraries/ghc-prim/GHC/Magic.hs b/libraries/ghc-prim/GHC/Magic.hs
index 1a6af9237d..22db69f2ae 100644
--- a/libraries/ghc-prim/GHC/Magic.hs
+++ b/libraries/ghc-prim/GHC/Magic.hs
@@ -69,6 +69,13 @@ lazy x = x
 -- argument will be called at most once, which may (or may not) enable certain
 -- optimizations. It can be useful to improve the performance of code in continuation
 -- passing style.
+--
+-- If 'oneShot' is used wrongly, then it may be that computations whose result
+-- 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 f = f
 -- Implementation note: This is wired in in MkId.lhs, so the code here is
diff --git a/testsuite/tests/typecheck/should_compile/T10744.hs b/testsuite/tests/typecheck/should_compile/T10744.hs
new file mode 100644
index 0000000000..64219ad696
--- /dev/null
+++ b/testsuite/tests/typecheck/should_compile/T10744.hs
@@ -0,0 +1,17 @@
+{-# LANGUAGE MagicHash #-}
+module T10744 where
+
+import GHC.Exts
+import GHC.Magic
+
+-- Checks if oneShot is open-kinded
+
+f0 :: Int -> Int
+f0 = oneShot $ \n -> n
+
+f1 :: Int# -> Int
+f1 = oneShot $ \n# -> I# n#
+
+f2 :: Int -> Int#
+f2 = oneShot $ \(I# n#) -> n#
+
diff --git a/testsuite/tests/typecheck/should_compile/all.T b/testsuite/tests/typecheck/should_compile/all.T
index 47a154a849..d9f2bd8160 100644
--- a/testsuite/tests/typecheck/should_compile/all.T
+++ b/testsuite/tests/typecheck/should_compile/all.T
@@ -468,3 +468,4 @@ test('T10562', normal, compile, [''])
 test('T10564', normal, compile, [''])
 test('T10632', normal, compile, [''])
 test('T10642', normal, compile, [''])
+test('T10744', normal, compile, [''])