StgLint: Lint constructor applications and strict workers for arity.

This will mean T9208 when run with lint will return a lint error instead
of resulting in a panic.

Fixes #21117

2 files changed, 33 insertions, 6 deletions

diff --git a/compiler/GHC/Stg/InferTags/Rewrite.hs b/compiler/GHC/Stg/InferTags/Rewrite.hs
index 1d2d280f2c..e35f700377 100644
--- a/compiler/GHC/Stg/InferTags/Rewrite.hs
+++ b/compiler/GHC/Stg/InferTags/Rewrite.hs
@@ -480,14 +480,20 @@ mkSeqs args untaggedIds mkExpr = do
 
 -- Out of all arguments passed at runtime only return these ending up in a
 -- strict field
-getStrictConArgs :: DataCon -> [a] -> [a]
+getStrictConArgs :: Outputable a => DataCon -> [a] -> [a]
 getStrictConArgs con args
     -- These are always lazy in their arguments.
     | isUnboxedTupleDataCon con = []
     | isUnboxedSumDataCon con = []
     -- For proper data cons we have to check.
     | otherwise =
+        assertPpr   (length args == length (dataConRuntimeRepStrictness con))
+                    (text "Missmatched con arg and con rep strictness lengths:" $$
+                     text "Con" <> ppr con <+> text "is applied to" <+> ppr args $$
+                     text "But seems to have arity" <> ppr (length repStrictness)) $
         [ arg | (arg,MarkedStrict)
                     <- zipEqual "getStrictConArgs"
                                 args
-                                (dataConRuntimeRepStrictness con)]
+                                repStrictness]
+        where
+            repStrictness = (dataConRuntimeRepStrictness con)
diff --git a/compiler/GHC/Stg/Lint.hs b/compiler/GHC/Stg/Lint.hs
index 657aa1603f..45e7b38471 100644
--- a/compiler/GHC/Stg/Lint.hs
+++ b/compiler/GHC/Stg/Lint.hs
@@ -47,7 +47,7 @@ import GHC.Core.DataCon
 import GHC.Core             ( AltCon(..) )
 import GHC.Core.Type
 
-import GHC.Types.Basic      ( TopLevelFlag(..), isTopLevel )
+import GHC.Types.Basic      ( TopLevelFlag(..), isTopLevel, isMarkedCbv )
 import GHC.Types.CostCentre ( isCurrentCCS )
 import GHC.Types.Error      ( DiagnosticReason(WarningWithoutFlag) )
 import GHC.Types.Id
@@ -69,6 +69,7 @@ import GHC.Data.Bag         ( Bag, emptyBag, isEmptyBag, snocBag, bagToList )
 import Control.Applicative ((<|>))
 import Control.Monad
 import Data.Maybe
+import GHC.Utils.Misc
 
 lintStgTopBindings :: forall a . (OutputablePass a, BinderP a ~ Id)
                    => Logger
@@ -179,10 +180,13 @@ lintStgRhs (StgRhsClosure _ _ _ binders expr)
         lintStgExpr expr
 
 lintStgRhs rhs@(StgRhsCon _ con _ _ args) = do
+    opts <- getStgPprOpts
     when (isUnboxedTupleDataCon con || isUnboxedSumDataCon con) $ do
-      opts <- getStgPprOpts
       addErrL (text "StgRhsCon is an unboxed tuple or sum application" $$
                pprStgRhs opts rhs)
+
+    lintConApp con args (pprStgRhs opts rhs)
+
     mapM_ lintStgArg args
     mapM_ checkPostUnariseConArg args
 
@@ -200,7 +204,7 @@ lintStgExpr e@(StgApp fun args) = do
       -- always needs to be applied to n arguments.
       -- See Note [Strict Worker Ids].
       let marks = fromMaybe [] $ idCbvMarks_maybe fun
-      if length marks > length args
+      if length (dropWhileEndLE (not . isMarkedCbv) marks) > length args
         then addErrL $ hang (text "Undersatured cbv marked ID in App" <+> ppr e ) 2 $
           (text "marks" <> ppr marks $$
           text "args" <> ppr args $$
@@ -211,10 +215,15 @@ lintStgExpr e@(StgApp fun args) = do
 lintStgExpr app@(StgConApp con _n args _arg_tys) = do
     -- unboxed sums should vanish during unarise
     lf <- getLintFlags
-    when (lf_unarised lf && isUnboxedSumDataCon con) $ do
+    let !unarised = lf_unarised lf
+    when (unarised && isUnboxedSumDataCon con) $ do
       opts <- getStgPprOpts
       addErrL (text "Unboxed sum after unarise:" $$
                pprStgExpr opts app)
+
+    opts <- getStgPprOpts
+    lintConApp con args (pprStgExpr opts app)
+
     mapM_ lintStgArg args
     mapM_ checkPostUnariseConArg args
 
@@ -262,6 +271,18 @@ lintAlt GenStgAlt{ alt_con   = DataAlt _
     mapM_ checkPostUnariseBndr bndrs
     addInScopeVars bndrs (lintStgExpr rhs)
 
+-- Post unarise check we apply constructors to the right number of args.
+-- This can be violated by invalid use of unsafeCoerce as showcased by test
+-- T9208
+lintConApp :: Foldable t => DataCon -> t a -> SDoc -> LintM ()
+lintConApp con args app = do
+    unarised <- lf_unarised <$> getLintFlags
+    when (unarised &&
+          not (isUnboxedTupleDataCon con) &&
+          length (dataConRuntimeRepStrictness con) /= length args) $ do
+      addErrL (text "Constructor applied to incorrect number of arguments:" $$
+               text "Application:" <> app)
+
 {-
 ************************************************************************
 *                                                                      *