Don't eta-expand unsaturated primopswip/T16846

Previously, as described in Note [Primop wrappers], `hasNoBinding` would
return False in the case of `PrimOpId`s. This would result in eta
expansion of unsaturated primop applications during CorePrep. Not only
did this expansion result in unnecessary allocations, but it also meant
lead to rather nasty inconsistencies between the CAFfy-ness
determinations made by TidyPgm and CorePrep.

This fixes #16846.

10 files changed, 134 insertions, 35 deletions

diff --git a/compiler/basicTypes/Id.hs b/compiler/basicTypes/Id.hs
index e2dfe925b1..4dceb4bc03 100644
--- a/compiler/basicTypes/Id.hs
+++ b/compiler/basicTypes/Id.hs
@@ -524,9 +524,17 @@ hasNoBinding :: Id -> Bool
 -- Data constructor workers used to be things of this kind, but
 -- they aren't any more.  Instead, we inject a binding for
 -- them at the CorePrep stage.
+--
+-- 'PrimOpId's also used to be of this kind. See Note [Primop wrappers] in PrimOp.hs.
+-- for the history of this.
+--
+-- Note that CorePrep currently eta expands things no-binding things and this
+-- can cause quite subtle bugs. See Note [Eta expansion of hasNoBinding things
+-- in CorePrep] in CorePrep for details.
+--
 -- EXCEPT: unboxed tuples, which definitely have no binding
 hasNoBinding id = case Var.idDetails id of
-                        PrimOpId _       -> True        -- See Note [Primop wrappers]
+                        PrimOpId _       -> False   -- See Note [Primop wrappers] in PrimOp.hs
                         FCallId _        -> True
                         DataConWorkId dc -> isUnboxedTupleCon dc || isUnboxedSumCon dc
                         _                -> isCompulsoryUnfolding (idUnfolding id)
@@ -570,19 +578,6 @@ The easiest way to do this is for hasNoBinding to return True of all things
 that have compulsory unfolding.  Some Ids with a compulsory unfolding also
 have a binding, but it does not harm to say they don't here, and its a very
 simple way to fix #14561.
-
-Note [Primop wrappers]
-~~~~~~~~~~~~~~~~~~~~~~
-Currently hasNoBinding claims that PrimOpIds don't have a curried
-function definition.  But actually they do, in GHC.PrimopWrappers,
-which is auto-generated from prelude/primops.txt.pp.  So actually, hasNoBinding
-could return 'False' for PrimOpIds.
-
-But we'd need to add something in CoreToStg to swizzle any unsaturated
-applications of GHC.Prim.plusInt# to GHC.PrimopWrappers.plusInt#.
-
-Nota Bene: GHC.PrimopWrappers is needed *regardless*, because it's
-used by GHCi, which does not implement primops direct at all.
 -}
 
 isDeadBinder :: Id -> Bool
diff --git a/compiler/basicTypes/Unique.hs b/compiler/basicTypes/Unique.hs
index adb23e0224..0031074a0b 100644
--- a/compiler/basicTypes/Unique.hs
+++ b/compiler/basicTypes/Unique.hs
@@ -46,7 +46,7 @@ module Unique (
         -- now all the built-in Uniques (and functions to make them)
         -- [the Oh-So-Wonderful Haskell module system wins again...]
         mkAlphaTyVarUnique,
-        mkPrimOpIdUnique,
+        mkPrimOpIdUnique, mkPrimOpWrapperUnique,
         mkPreludeMiscIdUnique, mkPreludeDataConUnique,
         mkPreludeTyConUnique, mkPreludeClassUnique,
         mkCoVarUnique,
@@ -368,6 +368,8 @@ mkPreludeClassUnique   :: Int -> Unique
 mkPreludeTyConUnique   :: Int -> Unique
 mkPreludeDataConUnique :: Arity -> Unique
 mkPrimOpIdUnique       :: Int -> Unique
+-- See Note [Primop wrappers] in PrimOp.hs.
+mkPrimOpWrapperUnique  :: Int -> Unique
 mkPreludeMiscIdUnique  :: Int -> Unique
 mkCoVarUnique          :: Int -> Unique
 
@@ -405,7 +407,8 @@ dataConWorkerUnique  u = incrUnique u
 dataConTyRepNameUnique u = stepUnique u 2
 
 --------------------------------------------------
-mkPrimOpIdUnique op         = mkUnique '9' op
+mkPrimOpIdUnique op         = mkUnique '9' (2*op)
+mkPrimOpWrapperUnique op    = mkUnique '9' (2*op+1)
 mkPreludeMiscIdUnique  i    = mkUnique '0' i
 
 -- The "tyvar uniques" print specially nicely: a, b, c, etc.
diff --git a/compiler/coreSyn/CorePrep.hs b/compiler/coreSyn/CorePrep.hs
index 98bffd3777..6be5346ab5 100644
--- a/compiler/coreSyn/CorePrep.hs
+++ b/compiler/coreSyn/CorePrep.hs
@@ -71,7 +71,7 @@ import qualified Data.Set as S
 
 The goal of this pass is to prepare for code generation.
 
-1.  Saturate constructor and primop applications.
+1.  Saturate constructor applications.
 
 2.  Convert to A-normal form; that is, function arguments
     are always variables.
@@ -1063,8 +1063,21 @@ because that has different strictness.  Hence the use of 'allLazy'.
 -- Building the saturated syntax
 -- ---------------------------------------------------------------------------
 
-maybeSaturate deals with saturating primops and constructors
-The type is the type of the entire application
+Note [Eta expansion of hasNoBinding things in CorePrep]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+maybeSaturate deals with eta expanding to saturate things that can't deal with
+unsaturated applications (identified by 'hasNoBinding', currently just
+foreign calls and unboxed tuple/sum constructors).
+
+Note that eta expansion in CorePrep is very fragile due to the "prediction" of
+CAFfyness made by TidyPgm (see Note [CAFfyness inconsistencies due to eta
+expansion in CorePrep] in TidyPgm for details.  We previously saturated primop
+applications here as well but due to this fragility (see #16846) we now deal
+with this another way, as described in Note [Primop wrappers] in PrimOp.
+
+It's quite likely that eta expansion of constructor applications will
+eventually break in a similar way to how primops did. We really should
+eliminate this case as well.
 -}
 
 maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
diff --git a/compiler/main/TidyPgm.hs b/compiler/main/TidyPgm.hs
index 6c5cf6f9f0..c0c6ffc3c3 100644
--- a/compiler/main/TidyPgm.hs
+++ b/compiler/main/TidyPgm.hs
@@ -1285,7 +1285,48 @@ So we have to *predict* the result here, which is revolting.
 
 In particular CorePrep expands Integer and Natural literals. So in the
 prediction code here we resort to applying the same expansion (cvt_literal).
-Ugh!
+There are also numberous other ways in which we can introduce inconsistencies
+between CorePrep and TidyPgm. See Note [CAFfyness inconsistencies due to eta
+expansion in TidyPgm] for one such example.
+
+Ugh! What ugliness we hath wrought.
+
+
+Note [CAFfyness inconsistencies due to eta expansion in TidyPgm]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Eta expansion during CorePrep can have non-obvious negative consequences on
+the CAFfyness computation done by TidyPgm (see Note [Disgusting computation of
+CafRefs] in TidyPgm). This late expansion happens/happened for a few reasons:
+
+ * CorePrep previously eta expanded unsaturated primop applications, as
+   described in Note [Primop wrappers]).
+
+ * CorePrep still does eta expand unsaturated data constructor applications.
+
+In particular, consider the program:
+
+    data Ty = Ty (RealWorld# -> (# RealWorld#, Int #))
+
+    -- Is this CAFfy?
+    x :: STM Int
+    x = Ty (retry# @Int)
+
+Consider whether x is CAFfy. One might be tempted to answer "no".
+Afterall, f obviously has no CAF references and the application (retry#
+@Int) is essentially just a variable reference at runtime.
+
+However, when CorePrep expanded the unsaturated application of 'retry#'
+it would rewrite this to
+
+    x = \u []
+       let sat = retry# @Int
+       in Ty sat
+
+This is now a CAF. Failing to handle this properly was the cause of
+#16846. We fixed this by eliminating the need to eta expand primops, as
+described in Note [Primop wrappers]), However we have not yet done the same for
+data constructor applications.
+
 -}
 
 type CafRefEnv = (VarEnv Id, LitNumType -> Integer -> Maybe CoreExpr)
diff --git a/compiler/prelude/PrelInfo.hs b/compiler/prelude/PrelInfo.hs
index 8ff9b19b45..204b7ce9f9 100644
--- a/compiler/prelude/PrelInfo.hs
+++ b/compiler/prelude/PrelInfo.hs
@@ -131,6 +131,7 @@ knownKeyNames
 
              , map idName wiredInIds
              , map (idName . primOpId) allThePrimOps
+             , map (idName . primOpWrapperId) allThePrimOps
              , basicKnownKeyNames
              , templateHaskellNames
              ]
diff --git a/compiler/prelude/PrelNames.hs b/compiler/prelude/PrelNames.hs
index c4956ad98b..2ed73d269a 100644
--- a/compiler/prelude/PrelNames.hs
+++ b/compiler/prelude/PrelNames.hs
@@ -498,7 +498,8 @@ pRELUDE :: Module
 pRELUDE         = mkBaseModule_ pRELUDE_NAME
 
 gHC_PRIM, gHC_TYPES, gHC_GENERICS, gHC_MAGIC,
-    gHC_CLASSES, gHC_BASE, gHC_ENUM, gHC_GHCI, gHC_GHCI_HELPERS, gHC_CSTRING,
+    gHC_CLASSES, gHC_PRIMOPWRAPPERS, gHC_BASE, gHC_ENUM,
+    gHC_GHCI, gHC_GHCI_HELPERS, gHC_CSTRING,
     gHC_SHOW, gHC_READ, gHC_NUM, gHC_MAYBE, gHC_INTEGER_TYPE, gHC_NATURAL,
     gHC_LIST, gHC_TUPLE, dATA_TUPLE, dATA_EITHER, dATA_STRING,
     dATA_FOLDABLE, dATA_TRAVERSABLE,
@@ -516,6 +517,7 @@ gHC_TYPES       = mkPrimModule (fsLit "GHC.Types")
 gHC_MAGIC       = mkPrimModule (fsLit "GHC.Magic")
 gHC_CSTRING     = mkPrimModule (fsLit "GHC.CString")
 gHC_CLASSES     = mkPrimModule (fsLit "GHC.Classes")
+gHC_PRIMOPWRAPPERS = mkPrimModule (fsLit "GHC.PrimopWrappers")
 
 gHC_BASE        = mkBaseModule (fsLit "GHC.Base")
 gHC_ENUM        = mkBaseModule (fsLit "GHC.Enum")
diff --git a/compiler/prelude/PrimOp.hs b/compiler/prelude/PrimOp.hs
index 3e157aea9b..ac4f162e08 100644
--- a/compiler/prelude/PrimOp.hs
+++ b/compiler/prelude/PrimOp.hs
@@ -13,6 +13,7 @@ module PrimOp (
         PrimOp(..), PrimOpVecCat(..), allThePrimOps,
         primOpType, primOpSig,
         primOpTag, maxPrimOpTag, primOpOcc,
+        primOpWrapperId,
 
         tagToEnumKey,
 
@@ -34,14 +35,18 @@ import TysWiredIn
 
 import CmmType
 import Demand
-import OccName          ( OccName, pprOccName, mkVarOccFS )
+import Id               ( Id, mkVanillaGlobalWithInfo )
+import IdInfo           ( vanillaIdInfo, setCafInfo, CafInfo(NoCafRefs) )
+import Name
+import PrelNames        ( gHC_PRIMOPWRAPPERS )
 import TyCon            ( TyCon, isPrimTyCon, PrimRep(..) )
 import Type
 import RepType          ( typePrimRep1, tyConPrimRep1 )
 import BasicTypes       ( Arity, Fixity(..), FixityDirection(..), Boxity(..),
                           SourceText(..) )
+import SrcLoc           ( wiredInSrcSpan )
 import ForeignCall      ( CLabelString )
-import Unique           ( Unique, mkPrimOpIdUnique )
+import Unique           ( Unique, mkPrimOpIdUnique, mkPrimOpWrapperUnique )
 import Outputable
 import FastString
 import Module           ( UnitId )
@@ -572,6 +577,47 @@ primOpOcc op = case primOpInfo op of
                Compare   occ _     -> occ
                GenPrimOp occ _ _ _ -> occ
 
+{- Note [Primop wrappers]
+~~~~~~~~~~~~~~~~~~~~~~~~~
+Previously hasNoBinding would claim that PrimOpIds didn't have a curried
+function definition. This caused quite some trouble as we would be forced to
+eta expand unsaturated primop applications very late in the Core pipeline. Not
+only would this produce unnecessary thunks, but it would also result in nasty
+inconsistencies in CAFfy-ness determinations (see #16846 and
+Note [CAFfyness inconsistencies due to late eta expansion] in TidyPgm).
+
+However, it was quite unnecessary for hasNoBinding to claim this; primops in
+fact *do* have curried definitions which are found in GHC.PrimopWrappers, which
+is auto-generated by utils/genprimops from prelude/primops.txt.pp. These wrappers
+are standard Haskell functions mirroring the types of the primops they wrap.
+For instance, in the case of plusInt# we would have:
+
+    module GHC.PrimopWrappers where
+    import GHC.Prim as P
+    plusInt# a b = P.plusInt# a b
+
+We now take advantage of these curried definitions by letting hasNoBinding
+claim that PrimOpIds have a curried definition and then rewrite any unsaturated
+PrimOpId applications that we find during CoreToStg as applications of the
+associated wrapper (e.g. `GHC.Prim.plusInt# 3#` will get rewritten to
+`GHC.PrimopWrappers.plusInt# 3#`).` The Id of the wrapper for a primop can be
+found using 'PrimOp.primOpWrapperId'.
+
+Nota Bene: GHC.PrimopWrappers is needed *regardless*, because it's
+used by GHCi, which does not implement primops direct at all.
+
+-}
+
+-- | Returns the 'Id' of the wrapper associated with the given 'PrimOp'.
+-- See Note [Primop wrappers].
+primOpWrapperId :: PrimOp -> Id
+primOpWrapperId op = mkVanillaGlobalWithInfo name ty info
+  where
+    info = setCafInfo vanillaIdInfo NoCafRefs
+    name = mkExternalName uniq gHC_PRIMOPWRAPPERS (primOpOcc op) wiredInSrcSpan
+    uniq = mkPrimOpWrapperUnique (primOpTag op)
+    ty   = primOpType op
+
 isComparisonPrimOp :: PrimOp -> Bool
 isComparisonPrimOp op = case primOpInfo op of
                           Compare {} -> True
diff --git a/compiler/stgSyn/CoreToStg.hs b/compiler/stgSyn/CoreToStg.hs
index caa012124b..6c59ebb081 100644
--- a/compiler/stgSyn/CoreToStg.hs
+++ b/compiler/stgSyn/CoreToStg.hs
@@ -45,7 +45,7 @@ import Util
 import DynFlags
 import ForeignCall
 import Demand           ( isUsedOnce )
-import PrimOp           ( PrimCall(..) )
+import PrimOp           ( PrimCall(..), primOpWrapperId )
 import SrcLoc           ( mkGeneralSrcSpan )
 
 import Data.List.NonEmpty (nonEmpty, toList)
@@ -537,8 +537,12 @@ coreToStgApp _ f args ticks = do
                                       (dropRuntimeRepArgs (fromMaybe [] (tyConAppArgs_maybe res_ty)))
 
                 -- Some primitive operator that might be implemented as a library call.
-                PrimOpId op      -> ASSERT( saturated )
-                                    StgOpApp (StgPrimOp op) args' res_ty
+                -- As described in Note [Primop wrappers] in PrimOp.hs, here we
+                -- turn unsaturated primop applications into applications of
+                -- the primop's wrapper.
+                PrimOpId op
+                  | saturated    -> StgOpApp (StgPrimOp op) args' res_ty
+                  | otherwise    -> StgApp (primOpWrapperId op) args'
 
                 -- A call to some primitive Cmm function.
                 FCallId (CCall (CCallSpec (StaticTarget _ lbl (Just pkgId) True)
diff --git a/testsuite/tests/codeGen/should_fail/T13233.hs b/testsuite/tests/codeGen/should_fail/T13233.hs
index f24fc03bfb..a8d2343e65 100644
--- a/testsuite/tests/codeGen/should_fail/T13233.hs
+++ b/testsuite/tests/codeGen/should_fail/T13233.hs
@@ -21,6 +21,9 @@ obscure _ = ()
 quux :: ()
 quux = obscure (#,#)
 
+-- It used to be that primops has no binding. However, as described in
+-- Note [Primop wrappers] in PrimOp we now rewrite unsaturated primop
+-- applications to their wrapper, which allows safe use of levity polymorphism.
 primop :: forall (rep :: RuntimeRep) (a :: TYPE rep) b c.
           a -> b -> (State# RealWorld -> (# State# RealWorld, c #))
        -> State# RealWorld -> (# State# RealWorld, Weak# b #)
diff --git a/testsuite/tests/codeGen/should_fail/T13233.stderr b/testsuite/tests/codeGen/should_fail/T13233.stderr
index 1531abed8e..08f1f62a88 100644
--- a/testsuite/tests/codeGen/should_fail/T13233.stderr
+++ b/testsuite/tests/codeGen/should_fail/T13233.stderr
@@ -14,12 +14,3 @@ T13233.hs:22:16: error:
     Levity-polymorphic arguments:
       a :: TYPE rep1
       b :: TYPE rep2
-
-T13233.hs:27:10: error:
-    Cannot use function with levity-polymorphic arguments:
-      mkWeak# :: a
-                 -> b
-                 -> (State# RealWorld -> (# State# RealWorld, c #))
-                 -> State# RealWorld
-                 -> (# State# RealWorld, Weak# b #)
-    Levity-polymorphic arguments: a :: TYPE rep