From c40cbaa22973faf38599625351f9e2dd8dac3691 Mon Sep 17 00:00:00 2001
From: Andreas Klebinger <klebinger.andreas@gmx.at>
Date: Mon, 7 Mar 2022 21:42:19 +0100
Subject: Improve -dtag-inference-checks checks.

FUN closures don't get tagged when evaluated. So no point in checking their
tags.
---
 compiler/GHC/Stg/InferTags.hs     |  2 ++
 compiler/GHC/StgToCmm/DataCon.hs  |  3 ++-
 compiler/GHC/StgToCmm/TagCheck.hs | 48 +++++++++++++++++++++++++++++----------
 3 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/compiler/GHC/Stg/InferTags.hs b/compiler/GHC/Stg/InferTags.hs
index 79acaa14f5..6d28f447d5 100644
--- a/compiler/GHC/Stg/InferTags.hs
+++ b/compiler/GHC/Stg/InferTags.hs
@@ -186,6 +186,8 @@ As it makes little difference for runtime performance I've treated functions as
 it made the code simpler. But besides implementation complexity there isn't any reason
 why we couldn't be more rigourous in dealing with functions.
 
+NB: It turned out because of #21193 option two wouldn't really have been an option anyway.
+
 Note [Tag inference debugging]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 There is a flag -dtag-inference-checks which inserts various
diff --git a/compiler/GHC/StgToCmm/DataCon.hs b/compiler/GHC/StgToCmm/DataCon.hs
index 2a7203e101..89bdb88058 100644
--- a/compiler/GHC/StgToCmm/DataCon.hs
+++ b/compiler/GHC/StgToCmm/DataCon.hs
@@ -212,7 +212,8 @@ buildDynCon' binder mn actually_bound ccs con args
           ; let ticky_name | actually_bound = Just binder
                            | otherwise = Nothing
 
-          ; checkConArgsDyn (text "TagCheck failed - con_alloc:" <> ppr binder) con (map fromNonVoid args)
+          ; checkConArgsDyn (hang (text "TagCheck failed on constructor application.") 4 $
+                                   text "On binder:" <> ppr binder $$ text "Constructor:" <> ppr con) con (map fromNonVoid args)
           ; hp_plus_n <- allocDynClosure ticky_name info_tbl lf_info
                                           use_cc blame_cc args_w_offsets
           ; return (mkRhsInit platform reg lf_info hp_plus_n) }
diff --git a/compiler/GHC/StgToCmm/TagCheck.hs b/compiler/GHC/StgToCmm/TagCheck.hs
index 4c9fb23cc5..26a207c814 100644
--- a/compiler/GHC/StgToCmm/TagCheck.hs
+++ b/compiler/GHC/StgToCmm/TagCheck.hs
@@ -35,11 +35,13 @@ import GHC.Utils.Panic (pprPanic)
 import GHC.Utils.Panic.Plain (panic)
 import GHC.Stg.Syntax
 import GHC.StgToCmm.Closure
+import GHC.Cmm.Switch (mkSwitchTargets)
+import GHC.Cmm.Info (cmmGetClosureType)
 import GHC.Types.RepType (dataConRuntimeRepStrictness)
 import GHC.Types.Basic
 import GHC.Data.FastString (mkFastString)
-import GHC.Cmm.Info (cmmGetClosureType)
-import GHC.Cmm.Utils (mkWordCLit)
+
+import qualified Data.Map as M
 
 -- | Check all arguments marked as already tagged for a function
 -- are tagged by inserting runtime checks.
@@ -91,21 +93,43 @@ emitTagAssertion onWhat fun = do
   -- If there is no tag check if we are dealing with a PAP
   ; emitLabel lno_tag
   ; emitComment (mkFastString "closereTypeCheck")
-  ; align_check <- stgToCmmAlignCheck <$> getStgToCmmConfig
-  ; profile <- getProfile
-  ; let closure_ty = cmmGetClosureType profile align_check fun
-  ; ty_reg <- newTemp (bWord platform)
-  ; emitAssign (CmmLocal ty_reg) closure_ty
-  ; emit $ mkCbranch (cmmEqWord platform
-                        (CmmReg $ CmmLocal ty_reg)
-                        (CmmLit $ mkWordCLit platform PAP))
-                      lret lbarf (Just True)
+  ; needsArgTag fun lbarf lret
 
   ; emitLabel lbarf
   ; emitBarf ("Tag inference failed on:" ++ onWhat)
   ; emitLabel lret
   }
 
+-- | Jump to the first block if the argument closure is subject
+--   to tagging requirements. Otherwise jump to the 2nd one.
+needsArgTag :: CmmExpr -> BlockId -> BlockId -> FCode ()
+needsArgTag closure fail lpass = do
+  profile <- getProfile
+  align_check <- stgToCmmAlignCheck <$> getStgToCmmConfig
+  let clo_ty_e = cmmGetClosureType profile align_check closure
+  -- The ENTER macro doesn't evaluate FUN/PAP/BCO objects. So we
+  -- have to accept them not being tagged. See #21193
+  -- See Note [TagInfo of functions]
+  let targets = mkSwitchTargets
+        False
+        (INVALID_OBJECT, N_CLOSURE_TYPES)
+        (Just fail)
+        (M.fromList [(PAP,lpass)
+                    ,(BCO,lpass)
+                    ,(FUN,lpass)
+                    ,(FUN_1_0,lpass)
+                    ,(FUN_0_1,lpass)
+                    ,(FUN_2_0,lpass)
+                    ,(FUN_1_1,lpass)
+                    ,(FUN_0_2,lpass)
+                    ,(FUN_STATIC,lpass)
+                    ])
+
+  emit $ mkSwitch clo_ty_e targets
+
+  emit $ mkBranch lpass
+
+
 emitArgTagCheck :: SDoc -> [CbvMark] -> [Id] -> FCode ()
 emitArgTagCheck info marks args = whenCheckTags $ do
   mod <- getModuleName
@@ -138,7 +162,7 @@ checkArg msg MarkedCbv arg = whenCheckTags $
       if taggedCgInfo info
           then return ()
           else case (cg_loc info) of
-            CmmLoc loc -> emitTagAssertion (showPprUnsafe msg) loc
+            CmmLoc loc -> emitTagAssertion (showPprUnsafe $ msg <+> text "arg:" <> ppr arg) loc
             LneLoc {} -> panic "LNE-arg"
 
 -- Check that argument is properly tagged.
-- 
cgit v1.2.1