Introduce keepAlive primop

4 files changed, 54 insertions, 1 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index 2ba94c1982..454c2e3a9e 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -3121,6 +3121,20 @@ primop NumSparks "numSparks#" GenPrimOp
    has_side_effects = True
    out_of_line = True
 
+
+------------------------------------------------------------------------
+section "Controlling object lifetime"
+        {Ensuring that objects don't die a premature death.}
+------------------------------------------------------------------------
+
+-- See Note [keepAlive# magic] in GHC.CoreToStg.Prep.
+primop KeepAliveOp "keepAlive#" GenPrimOp
+   o -> State# RealWorld -> (State# RealWorld -> p) -> p
+   { TODO. }
+   with
+   strictness = { \ _arity -> mkClosedStrictSig [topDmd, topDmd, strictOnceApply1Dmd] topDiv }
+
+
 ------------------------------------------------------------------------
 section "Tag to enum stuff"
         {Convert back and forth between values of enumerated types
diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index a67955ad2b..1e34a5fd62 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -1642,6 +1642,8 @@ app_ok primop_ok fun args
         -> False       --     for the special cases for SeqOp and DataToTagOp
         | DataToTagOp <- op
         -> False
+        | KeepAliveOp <- op
+        -> False
 
         | otherwise
         -> primop_ok op  -- Check the primop itself
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index 3b3921f5e2..01934423ed 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -32,7 +32,10 @@ import GHC.Tc.Utils.Env
 import GHC.Unit
 
 import GHC.Builtin.Names
+import GHC.Builtin.PrimOps
 import GHC.Builtin.Types
+import GHC.Builtin.Types.Prim ( realWorldStatePrimTy )
+import GHC.Types.Id.Make ( realWorldPrimId, mkPrimOpId )
 
 import GHC.Core.Utils
 import GHC.Core.Opt.Arity
@@ -47,6 +50,7 @@ import GHC.Core.TyCon
 import GHC.Core.DataCon
 import GHC.Core.Opt.OccurAnal
 
+
 import GHC.Data.Maybe
 import GHC.Data.OrdList
 import GHC.Data.FastString
@@ -64,7 +68,6 @@ import GHC.Types.Var.Set
 import GHC.Types.Var.Env
 import GHC.Types.Id
 import GHC.Types.Id.Info
-import GHC.Types.Id.Make ( realWorldPrimId )
 import GHC.Types.Basic
 import GHC.Types.Name   ( NamedThing(..), nameSrcSpan, isInternalName )
 import GHC.Types.SrcLoc ( SrcSpan(..), realSrcLocSpan, mkRealSrcLoc )
@@ -790,6 +793,38 @@ cpeApp top_env expr
         -- rather than the far superior "f x y".  Test case is par01.
         = let (terminal, args', depth') = collect_args arg
           in cpe_app env terminal (args' ++ args) (depth + depth' - 1)
+
+    cpe_app env
+            (Var f)
+            args
+            n
+        | Just KeepAliveOp <- isPrimOpId_maybe f
+        , CpeApp (Type arg_rep)
+          : CpeApp (Type arg_ty)
+          : CpeApp (Type _result_rep)
+          : CpeApp (Type result_ty)
+          : CpeApp arg
+          : CpeApp s0
+          : CpeApp k
+          : rest <- pprTrace "cpe_app keepAlive#" (ppr args) args
+        = do { pprTraceM "cpe_app(keepAlive#)" (ppr n)
+             ; y <- newVar result_ty
+             ; s2 <- newVar realWorldStatePrimTy
+             ; -- beta reduce if possible
+             ; (floats, k') <- case k of
+                  Lam s body -> cpe_app (extendCorePrepEnvExpr env s s0) body rest (n-2)
+                  _          -> cpe_app env k (CpeApp s0 : rest) (n-1)
+             ; let touchId = mkPrimOpId TouchOp
+                   expr = Case k' y result_ty [Alt DEFAULT [] rhs]
+                   rhs = let scrut = mkApps (Var touchId) [Type arg_rep, Type arg_ty, arg, Var realWorldPrimId]
+                         in Case scrut s2 result_ty [Alt DEFAULT [] (Var y)]
+             ; pprTraceM "cpe_app(keepAlive)" (ppr expr)
+             ; (floats', expr') <- cpeBody env expr
+             ; return (floats `appendFloats` floats', expr')
+             }
+        | Just KeepAliveOp <- isPrimOpId_maybe f
+        = panic "invalid keepAlive# application"
+
     cpe_app env (Var f) (CpeApp _runtimeRep@Type{} : CpeApp _type@Type{} : CpeApp arg : rest) n
         | f `hasKey` runRWKey
         -- N.B. While it may appear that n == 1 in the case of runRW#
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index b08edea624..2ea28a8eb2 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1623,6 +1623,8 @@ emitPrimOp dflags primop = case primop of
   TraceMarkerOp -> alwaysExternal
   SetThreadAllocationCounter -> alwaysExternal
 
+  KeepAliveOp -> panic "keepAlive# should have been eliminated in CorePrep"
+
  where
   profile = targetProfile dflags
   platform = profilePlatform profile