Refactor bad coercion checking in a few places

We do bad coercion checking in a few places in the compiler, but they
all checked it differently:

- CoreToStg.coreToStgArgs:

  Disallowed lifted-to-unlifted, disallowed changing prim reps even when
  the sizes are the same.

- StgCmmExpr.cgCase:

  Checked primRepSlot equality. This disallowed Int to Int64 coercions
  on 64-bit systems (and Int to Int32 on 32-bit) even though those are
  fine.

- CoreLint:

  Only place where we do this right. Full rules are explained in Note
  [Bad unsafe coercion].

This patch implements the check explained in Note [Bad unsafe coercion]
in CoreLint and uses it in CoreToStg.coreToStgArgs and
StgCmmExpr.cgCase.

This fixes #16952 and unblocks !1381 (which fixes #16893).

This is the most conservative and correct change I came up with that
fixes #16952.

One remaining problem with coercion checking is that it's currently done
in seemingly random places. What's special about CoreToStg.coreToStgArgs
and StgCmmExpr.cgCase? My guess is that adding assertions to those
places caught bugs before so we left assertions in those places. I think
we should remove these assertions and do coercion checking in CoreLint
and StgLint only (#17041).

4 files changed, 46 insertions, 34 deletions

diff --git a/compiler/GHC/StgToCmm/Expr.hs b/compiler/GHC/StgToCmm/Expr.hs
index 59cd246441..a8661d9de0 100644
--- a/compiler/GHC/StgToCmm/Expr.hs
+++ b/compiler/GHC/StgToCmm/Expr.hs
@@ -41,7 +41,7 @@ import Id
 import PrimOp
 import TyCon
 import Type             ( isUnliftedType )
-import RepType          ( isVoidTy, countConRepArgs, primRepSlot )
+import RepType          ( isVoidTy, countConRepArgs )
 import CostCentre       ( CostCentreStack, currentCCS )
 import Maybes
 import Util
@@ -50,7 +50,6 @@ import Outputable
 
 import Control.Monad (unless,void)
 import Control.Arrow (first)
-import Data.Function ( on )
 
 ------------------------------------------------------------------------
 --              cgExpr: the main function
@@ -428,10 +427,9 @@ assignment.
 -}
 cgCase (StgApp v []) bndr alt_type@(PrimAlt _) alts
   | isUnliftedType (idType v)  -- Note [Dodgy unsafeCoerce 1]
-  || reps_compatible
   = -- assignment suffices for unlifted types
     do { dflags <- getDynFlags
-       ; unless reps_compatible $
+       ; unless (reps_compatible dflags) $
            pprPanic "cgCase: reps do not match, perhaps a dodgy unsafeCoerce?"
                     (pp_bndr v $$ pp_bndr bndr)
        ; v_info <- getCgIdInfo v
@@ -441,13 +439,7 @@ cgCase (StgApp v []) bndr alt_type@(PrimAlt _) alts
        ; _ <- bindArgToReg (NonVoid bndr)
        ; cgAlts (NoGcInAlts,AssignedDirectly) (NonVoid bndr) alt_type alts }
   where
-    reps_compatible = ((==) `on` (primRepSlot . idPrimRep)) v bndr
-      -- Must compare SlotTys, not proper PrimReps, because with unboxed sums,
-      -- the types of the binders are generated from slotPrimRep and might not
-      -- match. Test case:
-      --   swap :: (# Int | Int #) -> (# Int | Int #)
-      --   swap (# x | #) = (# | x #)
-      --   swap (# | y #) = (# y | #)
+    reps_compatible dflags = primRepCompatible dflags (idPrimRep v) (idPrimRep bndr)
 
     pp_bndr id = ppr id <+> dcolon <+> ppr (idType id) <+> parens (ppr (idPrimRep id))
 
diff --git a/compiler/ghci/ByteCodeGen.hs b/compiler/ghci/ByteCodeGen.hs
index 2865aaeaa6..2ad089903b 100644
--- a/compiler/ghci/ByteCodeGen.hs
+++ b/compiler/ghci/ByteCodeGen.hs
@@ -1219,7 +1219,7 @@ generateCCall d0 s p (CCallSpec target cconv safety) fn args_r_to_l
          push_args    = concatOL pushs_arg
          !d_after_args = d0 + wordsToBytes dflags a_reps_sizeW
          a_reps_pushed_RAW
-            | null a_reps_pushed_r_to_l || head a_reps_pushed_r_to_l /= VoidRep
+            | null a_reps_pushed_r_to_l || not (isVoidRep (head a_reps_pushed_r_to_l))
             = panic "ByteCodeGen.generateCCall: missing or invalid World token?"
             | otherwise
             = reverse (tail a_reps_pushed_r_to_l)
@@ -1904,7 +1904,8 @@ atomPrimRep (AnnLit l)              = typePrimRep1 (literalType l)
 -- #12128:
 -- A case expression can be an atom because empty cases evaluate to bottom.
 -- See Note [Empty case alternatives] in coreSyn/CoreSyn.hs
-atomPrimRep (AnnCase _ _ ty _)      = ASSERT(typePrimRep ty == [LiftedRep]) LiftedRep
+atomPrimRep (AnnCase _ _ ty _)      =
+  ASSERT(case typePrimRep ty of [LiftedRep] -> True; _ -> False) LiftedRep
 atomPrimRep (AnnCoercion {})        = VoidRep
 atomPrimRep other = pprPanic "atomPrimRep" (ppr (deAnnotate' other))
 
diff --git a/compiler/stgSyn/CoreToStg.hs b/compiler/stgSyn/CoreToStg.hs
index d548616dde..4982ad68f4 100644
--- a/compiler/stgSyn/CoreToStg.hs
+++ b/compiler/stgSyn/CoreToStg.hs
@@ -263,7 +263,7 @@ coreTopBindToStg dflags this_mod env ccs (NonRec id rhs)
         how_bound = LetBound TopLet $! manifestArity rhs
 
         (stg_rhs, ccs') =
-            initCts env $
+            initCts dflags env $
               coreToTopStgRhs dflags ccs this_mod (id,rhs)
 
         bind = StgTopLifted $ StgNonRec id stg_rhs
@@ -286,7 +286,7 @@ coreTopBindToStg dflags this_mod env ccs (Rec pairs)
 
         -- generate StgTopBindings and CAF cost centres created for CAFs
         (ccs', stg_rhss)
-          = initCts env' $ do
+          = initCts dflags env' $ do
                mapAccumLM (\ccs rhs -> do
                             (rhs', ccs') <-
                               coreToTopStgRhs dflags ccs this_mod rhs
@@ -598,16 +598,12 @@ coreToStgArgs (arg : args) = do         -- Non-type argument
         -- This matters particularly when the function is a primop
         -- or foreign call.
         -- Wanted: a better solution than this hacky warning
+
+    dflags <- getDynFlags
     let
-        arg_ty = exprType arg
-        stg_arg_ty = stgArgType stg_arg
-        bad_args = (isUnliftedType arg_ty && not (isUnliftedType stg_arg_ty))
-                || (typePrimRep arg_ty /= typePrimRep stg_arg_ty)
-        -- In GHCi we coerce an argument of type BCO# (unlifted) to HValue (lifted),
-        -- and pass it to a function expecting an HValue (arg_ty).  This is ok because
-        -- we can treat an unlifted value as lifted.  But the other way round
-        -- we complain.
-        -- We also want to check if a pointer is cast to a non-ptr etc
+        arg_rep = typePrimRep (exprType arg)
+        stg_arg_rep = typePrimRep (stgArgType stg_arg)
+        bad_args = not (primRepsCompatible dflags arg_rep stg_arg_rep)
 
     WARN( bad_args, text "Dangerous-looking argument. Probable cause: bad unsafeCoerce#" $$ ppr arg )
      return (stg_arg : stg_args, ticks ++ aticks)
@@ -816,7 +812,8 @@ isPAP env _               = False
 -- *down*.
 
 newtype CtsM a = CtsM
-    { unCtsM :: IdEnv HowBound
+    { unCtsM :: DynFlags -- Needed for checking for bad coercions in coreToStgArgs
+             -> IdEnv HowBound
              -> a
     }
     deriving (Functor)
@@ -853,8 +850,8 @@ data LetInfo
 
 -- The std monad functions:
 
-initCts :: IdEnv HowBound -> CtsM a -> a
-initCts env m = unCtsM m env
+initCts :: DynFlags -> IdEnv HowBound -> CtsM a -> a
+initCts dflags env m = unCtsM m dflags env
 
 
 
@@ -862,11 +859,11 @@ initCts env m = unCtsM m env
 {-# INLINE returnCts #-}
 
 returnCts :: a -> CtsM a
-returnCts e = CtsM $ \_ -> e
+returnCts e = CtsM $ \_ _ -> e
 
 thenCts :: CtsM a -> (a -> CtsM b) -> CtsM b
-thenCts m k = CtsM $ \env
-  -> unCtsM (k (unCtsM m env)) env
+thenCts m k = CtsM $ \dflags env
+  -> unCtsM (k (unCtsM m dflags env)) dflags env
 
 instance Applicative CtsM where
     pure = returnCts
@@ -875,15 +872,18 @@ instance Applicative CtsM where
 instance Monad CtsM where
     (>>=)  = thenCts
 
+instance HasDynFlags CtsM where
+    getDynFlags = CtsM $ \dflags _ -> dflags
+
 -- Functions specific to this monad:
 
 extendVarEnvCts :: [(Id, HowBound)] -> CtsM a -> CtsM a
 extendVarEnvCts ids_w_howbound expr
-   =    CtsM $   \env
-   -> unCtsM expr (extendVarEnvList env ids_w_howbound)
+   =    CtsM $   \dflags env
+   -> unCtsM expr dflags (extendVarEnvList env ids_w_howbound)
 
 lookupVarCts :: Id -> CtsM HowBound
-lookupVarCts v = CtsM $ \env -> lookupBinding env v
+lookupVarCts v = CtsM $ \_ env -> lookupBinding env v
 
 lookupBinding :: IdEnv HowBound -> Id -> HowBound
 lookupBinding env v = case lookupVarEnv env v of
diff --git a/compiler/types/TyCon.hs b/compiler/types/TyCon.hs
index d3977aaecb..45fdb411ab 100644
--- a/compiler/types/TyCon.hs
+++ b/compiler/types/TyCon.hs
@@ -121,6 +121,8 @@ module TyCon(
         primRepSizeB,
         primElemRepSizeB,
         primRepIsFloat,
+        primRepsCompatible,
+        primRepCompatible,
 
         -- * Recursion breaking
         RecTcChecker, initRecTc, defaultRecTcMaxBound,
@@ -1422,7 +1424,7 @@ data PrimRep
   | FloatRep
   | DoubleRep
   | VecRep Int PrimElemRep  -- ^ A vector
-  deriving( Eq, Show )
+  deriving( Show )
 
 data PrimElemRep
   = Int8ElemRep
@@ -1452,6 +1454,23 @@ isGcPtrRep LiftedRep   = True
 isGcPtrRep UnliftedRep = True
 isGcPtrRep _           = False
 
+-- A PrimRep is compatible with another iff one can be coerced to the other.
+-- See Note [bad unsafe coercion] in CoreLint for when are two types coercible.
+primRepCompatible :: DynFlags -> PrimRep -> PrimRep -> Bool
+primRepCompatible dflags rep1 rep2 =
+    (isUnboxed rep1 == isUnboxed rep2) &&
+    (primRepSizeB dflags rep1 == primRepSizeB dflags rep2) &&
+    (primRepIsFloat rep1 == primRepIsFloat rep2)
+  where
+    isUnboxed = not . isGcPtrRep
+
+-- More general version of `primRepCompatible` for types represented by zero or
+-- more than one PrimReps.
+primRepsCompatible :: DynFlags -> [PrimRep] -> [PrimRep] -> Bool
+primRepsCompatible dflags reps1 reps2 =
+    length reps1 == length reps2 &&
+    and (zipWith (primRepCompatible dflags) reps1 reps2)
+
 -- | The size of a 'PrimRep' in bytes.
 --
 -- This applies also when used in a constructor, where we allow packing the