improve dead code elimination in CorePrep (fixes #7796)

2 files changed, 36 insertions, 74 deletions

diff --git a/compiler/coreSyn/CorePrep.lhs b/compiler/coreSyn/CorePrep.lhs
index 985f7f8c75..084c853382 100644
--- a/compiler/coreSyn/CorePrep.lhs
+++ b/compiler/coreSyn/CorePrep.lhs
@@ -13,6 +13,8 @@ module CorePrep (
 
 #include "HsVersions.h"
 
+import OccurAnal
+
 import HscTypes
 import PrelNames
 import CoreUtils
@@ -306,11 +308,12 @@ unreachable g$Bool and g$Unit functions.
 
 The way we fix this is to:
  * In cloneBndr, drop all unfoldings/rules
- * In deFloatTop, run a simple dead code analyser on each top-level RHS to drop
-   the dead local bindings. (we used to run the occurrence analyser to do
-   this job, but the occurrence analyser sometimes introduces new let
-   bindings for case binders, which lead to the bug in #5433, hence we
-   now have a special-purpose dead code analyser).
+
+ * In deFloatTop, run a simple dead code analyser on each top-level
+   RHS to drop the dead local bindings. For that call to OccAnal, we
+   disable the binder swap, else the occurrence analyser sometimes
+   introduces new let bindings for cased binders, which lead to the bug
+   in #5433.
 
 The reason we don't just OccAnal the whole output of CorePrep is that
 the tidier ensures that all top-level binders are GlobalIds, so they
@@ -1014,64 +1017,11 @@ deFloatTop (Floats _ floats)
     get b            _  = pprPanic "corePrepPgm" (ppr b)
 
     -- See Note [Dead code in CorePrep]
-    occurAnalyseRHSs (NonRec x e) = NonRec x (fst (dropDeadCode e))
-    occurAnalyseRHSs (Rec xes)    = Rec [ (x, fst (dropDeadCode e))
-                                        | (x, e) <- xes]
+    occurAnalyseRHSs (NonRec x e) = NonRec x (occurAnalyseExpr_NoBinderSwap e)
+    occurAnalyseRHSs (Rec xes)    = Rec [(x, occurAnalyseExpr_NoBinderSwap e) | (x, e) <- xes]
 
 ---------------------------------------------------------------------------
--- Simple dead-code analyser, see Note [Dead code in CorePrep]
-
-dropDeadCode :: CoreExpr -> (CoreExpr, VarSet)
-dropDeadCode (Var v)
-  = (Var v, if isLocalId v then unitVarSet v else emptyVarSet)
-dropDeadCode (App fun arg)
-  = (App fun' arg', fun_fvs `unionVarSet` arg_fvs)
-  where !(fun', fun_fvs) = dropDeadCode fun
-        !(arg', arg_fvs) = dropDeadCode arg
-dropDeadCode (Lam v e)
-  = (Lam v e', delVarSet fvs v)
-  where !(e', fvs) = dropDeadCode e
-dropDeadCode (Let (NonRec v rhs) body)
-  | v `elemVarSet` body_fvs
-  = (Let (NonRec v rhs') body', rhs_fvs `unionVarSet` (body_fvs `delVarSet` v))
-  | otherwise
-  = (body', body_fvs) -- drop the dead let bind!
-  where !(body', body_fvs) = dropDeadCode body
-        !(rhs',  rhs_fvs)  = dropDeadCode rhs
-dropDeadCode (Let (Rec prs) body)
-  | any (`elemVarSet` all_fvs) bndrs
-    -- approximation: strictly speaking we should do SCC analysis here,
-    -- but for simplicity we just look to see whether any of the binders
-    -- is used and drop the entire group if all are unused.
-  = (Let (Rec (zip bndrs rhss')) body', all_fvs `delVarSetList` bndrs)
-  | otherwise
-  = (body', body_fvs) -- drop the dead let bind!
-  where !(body', body_fvs) = dropDeadCode body
-        !(bndrs, rhss)     = unzip prs
-        !(rhss', rhs_fvss) = unzip (map dropDeadCode rhss)
-        all_fvs            = unionVarSets (body_fvs : rhs_fvss)
-
-dropDeadCode (Case scrut bndr t alts)
-  = (Case scrut' bndr t alts', scrut_fvs `unionVarSet` alts_fvs)
-  where !(scrut', scrut_fvs) = dropDeadCode scrut
-        !(alts',  alts_fvs)  = dropDeadCodeAlts alts
-dropDeadCode (Cast e c)
-  = (Cast e' c, fvs)
-  where !(e', fvs) = dropDeadCode e
-dropDeadCode (Tick t e)
-  = (Tick t e', fvs')
-  where !(e', fvs) = dropDeadCode e
-        fvs' | Breakpoint _ xs <- t =  fvs `unionVarSet` mkVarSet xs
-             | otherwise            =  fvs
-dropDeadCode e = (e, emptyVarSet)  -- Lit, Type, Coercion
-
-dropDeadCodeAlts :: [CoreAlt] -> ([CoreAlt], VarSet)
-dropDeadCodeAlts alts = (alts', unionVarSets fvss)
-  where !(alts', fvss) = unzip (map do_alt alts)
-        do_alt (c, vs, e) = ((c,vs,e'), fvs `delVarSetList` vs)
-          where !(e', fvs) = dropDeadCode e
-
--------------------------------------------
+
 canFloatFromNoCaf :: Platform -> Floats -> CpeRhs -> Maybe (Floats, CpeRhs)
        -- Note [CafInfo and floating]
 canFloatFromNoCaf platform (Floats ok_to_spec fs) rhs
diff --git a/compiler/simplCore/OccurAnal.lhs b/compiler/simplCore/OccurAnal.lhs
index dccbabcbdb..2c27070166 100644
--- a/compiler/simplCore/OccurAnal.lhs
+++ b/compiler/simplCore/OccurAnal.lhs
@@ -14,7 +14,7 @@ core expression with (hopefully) improved usage information.
 \begin{code}
 {-# LANGUAGE BangPatterns #-}
 module OccurAnal (
-        occurAnalysePgm, occurAnalyseExpr
+        occurAnalysePgm, occurAnalyseExpr, occurAnalyseExpr_NoBinderSwap
     ) where
 
 #include "HsVersions.h"
@@ -94,9 +94,16 @@ occurAnalysePgm this_mod active_rule imp_rules vects vectVars binds
 
 occurAnalyseExpr :: CoreExpr -> CoreExpr
         -- Do occurrence analysis, and discard occurence info returned
-occurAnalyseExpr expr
-  = snd (occAnal (initOccEnv all_active_rules) expr)
+occurAnalyseExpr = occurAnalyseExpr' True -- do binder swap
+
+occurAnalyseExpr_NoBinderSwap :: CoreExpr -> CoreExpr
+occurAnalyseExpr_NoBinderSwap = occurAnalyseExpr' False -- do not do binder swap
+
+occurAnalyseExpr' :: Bool -> CoreExpr -> CoreExpr
+occurAnalyseExpr' enable_binder_swap expr
+  = snd (occAnal env expr)
   where
+    env = (initOccEnv all_active_rules) {occ_binder_swap = enable_binder_swap}
     -- To be conservative, we say that all inlines and rules are active
     all_active_rules = \_ -> True
 \end{code}
@@ -1405,21 +1412,23 @@ occAnalAlt :: (OccEnv, Maybe (Id, CoreExpr))
 occAnalAlt (env, scrut_bind) case_bndr (con, bndrs, rhs)
   = case occAnal env rhs of { (rhs_usage1, rhs1) ->
     let
-        (rhs_usage2, rhs2) = wrapProxy scrut_bind case_bndr rhs_usage1 rhs1 
+        (rhs_usage2, rhs2) =
+          wrapProxy (occ_binder_swap env) scrut_bind case_bndr rhs_usage1 rhs1 
         (alt_usg, tagged_bndrs) = tagLamBinders rhs_usage2 bndrs
         bndrs' = tagged_bndrs      -- See Note [Binders in case alternatives]
     in
     (alt_usg, (con, bndrs', rhs2)) }
 
-wrapProxy :: Maybe (Id, CoreExpr) -> Id -> UsageDetails -> CoreExpr -> (UsageDetails, CoreExpr)
-wrapProxy (Just (scrut_var, rhs)) case_bndr body_usg body
-  | scrut_var `usedIn` body_usg
+wrapProxy :: Bool -> Maybe (Id, CoreExpr) -> Id -> UsageDetails -> CoreExpr -> (UsageDetails, CoreExpr)
+wrapProxy enable_binder_swap (Just (scrut_var, rhs)) case_bndr body_usg body
+  | enable_binder_swap,
+    scrut_var `usedIn` body_usg
   = ( body_usg' +++ unitVarEnv case_bndr NoOccInfo
     , Let (NonRec tagged_scrut_var rhs) body )
   where
     (body_usg', tagged_scrut_var) = tagBinder body_usg scrut_var
     
-wrapProxy _ _ body_usg body 
+wrapProxy _ _ _ body_usg body 
   = (body_usg, body)
 \end{code}
 
@@ -1432,11 +1441,13 @@ wrapProxy _ _ body_usg body
 
 \begin{code}
 data OccEnv
-  = OccEnv { occ_encl       :: !OccEncl      -- Enclosing context information
-           , occ_ctxt       :: !CtxtTy       -- Tells about linearity
-           , occ_gbl_scrut  :: GlobalScruts
-           , occ_rule_act   :: Activation -> Bool   -- Which rules are active
+  = OccEnv { occ_encl        :: !OccEncl      -- Enclosing context information
+           , occ_ctxt        :: !CtxtTy       -- Tells about linearity
+           , occ_gbl_scrut   :: GlobalScruts
+           , occ_rule_act    :: Activation -> Bool   -- Which rules are active
              -- See Note [Finding rule RHS free vars]
+           , occ_binder_swap :: !Bool -- enable the binder_swap
+             -- See CorePrep Note [Dead code in CorePrep]
     }
 
 type GlobalScruts = IdSet   -- See Note [Binder swap on GlobalId scrutinees]
@@ -1475,7 +1486,8 @@ initOccEnv active_rule
   = OccEnv { occ_encl  = OccVanilla
            , occ_ctxt  = []
            , occ_gbl_scrut = emptyVarSet --  PE emptyVarEnv emptyVarSet
-           , occ_rule_act = active_rule }
+           , occ_rule_act = active_rule
+           , occ_binder_swap = True }
 
 vanillaCtxt :: OccEnv -> OccEnv
 vanillaCtxt env = env { occ_encl = OccVanilla, occ_ctxt = [] }