[!14759] Allow Core plugins to access unoptimized Core (#23337)

compiler/GHC/Core/Opt/Pipeline.hs

@@ -13,6 +13,7 @@ import GHC.Prelude
 import GHC.Driver.DynFlags
 import GHC.Driver.Plugins ( withPlugins, installCoreToDos )
 import GHC.Driver.Env
+import GHC.Driver.Config (initSimpleOpts)
 import GHC.Driver.Config.Core.Lint ( endPass )
 import GHC.Driver.Config.Core.Opt.LiberateCase ( initLiberateCaseOpts )
 import GHC.Driver.Config.Core.Opt.Simplify ( initSimplifyOpts, initSimplMode, initGentleSimplMode )
@@ -21,9 +22,10 @@ import GHC.Driver.Config.Core.Rules ( initRuleOpts )
 import GHC.Platform.Ways  ( hasWay, Way(WayProf) )
 
 import GHC.Core
+import GHC.Core.SimpleOpt (simpleOptPgm)
 import GHC.Core.Opt.CSE  ( cseProgram )
 import GHC.Core.Rules   ( RuleBase, ruleCheckProgram, getRules )
-import GHC.Core.Ppr     ( pprCoreBindings )
+import GHC.Core.Ppr     ( pprCoreBindings, pprRules )
 import GHC.Core.Utils   ( dumpIdInfoOfProgram )
 import GHC.Core.Lint    ( lintAnnots )
 import GHC.Core.Lint.Interactive ( interactiveInScope )
@@ -202,10 +204,14 @@ getCoreToDo dflags hpt_rule_base extra_vars
 
     core_todo =
      [
-    -- We want to do the static argument transform before full laziness as it
-    -- may expose extra opportunities to float things outwards. However, to fix
-    -- up the output of the transformation we need at do at least one simplify
-    -- after this before anything else
+        -- We always perform a run of the simple optimizer after desugaring to
+        -- remove really bad code
+        CoreDesugarOpt,
+
+        -- We want to do the static argument transform before full laziness as it
+        -- may expose extra opportunities to float things outwards. However, to fix
+        -- up the output of the transformation we need at do at least one simplify
+        -- after this before anything else
         runWhen static_args (CoreDoPasses [ simpl_gently, CoreDoStaticArgs ]),
 
         -- initial simplify: mk specialiser happy: minimum effort please
@@ -467,6 +473,7 @@ doCorePass pass guts = do
   let fam_envs = (p_fam_env, mg_fam_inst_env guts)
   let updateBinds  f = return $ guts { mg_binds = f (mg_binds guts) }
   let updateBindsM f = f (mg_binds guts) >>= \b' -> return $ guts { mg_binds = b' }
+  let updateBindsAndRulesM f = f (mg_binds guts) (mg_rules guts) >>= \(b',r') -> return $ guts { mg_binds = b', mg_rules = r' }
   -- Important to force this now as name_ppr_ctx lives through an entire phase in
   -- the optimiser and if it's not forced then the entire previous `ModGuts` will
   -- be retained until the end of the phase. (See #24328 for more analysis)
@@ -479,6 +486,9 @@ doCorePass pass guts = do
 
 
   case pass of
+    CoreDesugarOpt            -> {-# SCC "DesugarOpt" #-}
+                                 updateBindsAndRulesM (desugarOpt dflags logger (mg_module guts))
+
     CoreDoSimplify opts       -> {-# SCC "Simplify" #-}
                                  liftIOWithCount $ simplifyPgm logger (hsc_unit_env hsc_env) name_ppr_ctx opts guts
 
@@ -537,7 +547,6 @@ doCorePass pass guts = do
     CoreDoPluginPass _ p      -> {-# SCC "Plugin" #-} p guts
 
     CoreDesugar               -> pprPanic "doCorePass" (ppr pass)
-    CoreDesugarOpt            -> pprPanic "doCorePass" (ppr pass)
     CoreTidy                  -> pprPanic "doCorePass" (ppr pass)
     CorePrep                  -> pprPanic "doCorePass" (ppr pass)
 
@@ -580,3 +589,25 @@ dmdAnal logger before_ww dflags fam_envs rules binds = do
     dumpIdInfoOfProgram (hasPprDebug dflags) (ppr . zapDmdEnvSig . dmdSigInfo) binds_plus_dmds
   -- See Note [Stamp out space leaks in demand analysis] in GHC.Core.Opt.DmdAnal
   seqBinds binds_plus_dmds `seq` return binds_plus_dmds
+
+
+-- | Simple optimization after desugaring.
+--
+-- This is used to remove the bad code that the desugarer produces (top-level
+-- dictionnary bindings, type bindings, etc.).
+--
+-- It does things that the real Simplifier doesn't do: e.g. floating-in
+-- top-level String literals. Hence we can't fully remove it.
+--
+-- It has been moved from being called by the desugarer directly to being the
+-- first Core-to-Core pass to accomodate Core plugins that want to see Core even
+-- before the first (simple) optimization took place. See #23337
+desugarOpt :: DynFlags -> Logger -> Module -> CoreProgram -> [CoreRule] -> CoreM (CoreProgram,[CoreRule])
+desugarOpt dflags logger mod binds rules = liftIO $ do
+  let simpl_opts = initSimpleOpts dflags
+  let !(ds_binds, ds_rules_for_imps, occ_anald_binds) = simpleOptPgm simpl_opts mod binds rules
+
+  putDumpFileMaybe logger Opt_D_dump_occur_anal "Occurrence analysis"
+    FormatCore (pprCoreBindings occ_anald_binds $$ pprRules ds_rules_for_imps )
+
+  pure (ds_binds, ds_rules_for_imps)

compiler/GHC/Core/Opt/Pipeline/Types.hs

@@ -58,8 +58,7 @@ data CoreToDo           -- These are diff core-to-core passes,
   | CoreDoPasses [CoreToDo]      -- lists of these things
 
   | CoreDesugar    -- Right after desugaring, no simple optimisation yet!
-  | CoreDesugarOpt -- CoreDesugarXXX: Not strictly a core-to-core pass, but produces
-                       --                 Core output, and hence useful to pass to endPass
+  | CoreDesugarOpt -- Simple optimisation after desugaring
 
   | CoreTidy
   | CorePrep

compiler/GHC/HsToCore.hs

@@ -48,7 +48,7 @@ import GHC.Core.TyCo.Compare( eqType )
 import GHC.Core.TyCon       ( tyConDataCons )
 import GHC.Core
 import GHC.Core.FVs       ( exprsSomeFreeVarsList, exprFreeVars )
-import GHC.Core.SimpleOpt ( simpleOptPgm, simpleOptExpr )
+import GHC.Core.SimpleOpt ( simpleOptExpr )
 import GHC.Core.Utils
 import GHC.Core.Unfold.Make
 import GHC.Core.Coercion
@@ -200,27 +200,18 @@ deSugar hsc_env
 
      do {       -- Add export flags to bindings
           keep_alive <- readIORef keep_var
-        ; let (rules_for_locals, rules_for_imps) = partition isLocalRule all_rules
+        ; let (rules_for_locals, ds_rules_for_imps) = partition isLocalRule all_rules
               final_prs = addExportFlagsAndRules bcknd export_set keep_alive
                                                  rules_for_locals (fromOL all_prs)
 
-              final_pgm = combineEvBinds ds_ev_binds final_prs
+              ds_binds = combineEvBinds ds_ev_binds final_prs
         -- Notice that we put the whole lot in a big Rec, even the foreign binds
         -- When compiling PrelFloat, which defines data Float = F# Float#
         -- we want F# to be in scope in the foreign marshalling code!
         -- You might think it doesn't matter, but the simplifier brings all top-level
         -- things into the in-scope set before simplifying; so we get no unfolding for F#!
 
-        ; endPassHscEnvIO hsc_env name_ppr_ctx CoreDesugar final_pgm rules_for_imps
-        ; let simpl_opts = initSimpleOpts dflags
-        ; let (ds_binds, ds_rules_for_imps, occ_anald_binds)
-                = simpleOptPgm simpl_opts mod final_pgm rules_for_imps
-                         -- The simpleOptPgm gets rid of type
-                         -- bindings plus any stupid dead code
-        ; putDumpFileMaybe logger Opt_D_dump_occur_anal "Occurrence analysis"
-            FormatCore (pprCoreBindings occ_anald_binds $$ pprRules ds_rules_for_imps )
-
-        ; endPassHscEnvIO hsc_env name_ppr_ctx CoreDesugarOpt ds_binds ds_rules_for_imps
+        ; endPassHscEnvIO hsc_env name_ppr_ctx CoreDesugar ds_binds ds_rules_for_imps
 
         ; let pluginModules = map lpModule (loadedPlugins (hsc_plugins hsc_env))
               home_unit = hsc_home_unit hsc_env

testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs

@@ -19,13 +19,14 @@ pass :: ModGuts -> CoreM ModGuts
 pass g = do
           dflags <- getDynFlags
           mapM_ (printAnn dflags g) (mg_binds g) >> return g
-  where printAnn :: DynFlags -> ModGuts -> CoreBind -> CoreM CoreBind
-        printAnn dflags guts bndr@(NonRec b _) = do
+  where printAnn :: DynFlags -> ModGuts -> CoreBind -> CoreM ()
+        printAnn dflags guts (NonRec b _) = lookupAnn dflags guts b
+        printAnn dflags guts (Rec ps) = mapM_ (lookupAnn dflags guts . fst) ps
+
+        lookupAnn dflags guts b = do
           anns <- annotationsOn guts b :: CoreM [SomeAnn]
           unless (null anns) $ putMsgS $
             "Annotated binding found: " ++  showSDoc dflags (ppr b)
-          return bndr
-        printAnn _ _ bndr = return bndr
 
 annotationsOn :: Data a => ModGuts -> CoreBndr -> CoreM [a]
 annotationsOn guts bndr = do

testsuite/tests/plugins/late-plugin/LatePlugin.hs

@@ -43,8 +43,13 @@ editCoreBinding early modName pgm = do
     pure $ go pgm
   where
     go :: [CoreBind] -> [CoreBind]
-    go (b@(NonRec v e) : bs)
-      | occNameString (getOccName v) == "testBinding" && exprType e `eqType` intTy =
-          NonRec v (mkUncheckedIntExpr $ bool 222222 111111 early) : bs
-    go (b:bs) = b : go bs
+    go (Rec ps : bs) = Rec (map (uncurry (go_bind (,))) ps) : go bs
+    go (NonRec v e : bs) = go_bind NonRec v e : go bs
     go [] = []
+
+    go_bind c v e
+      | occNameString (getOccName v) == "testBinding"
+      , exprType e `eqType` intTy
+      = c v (mkUncheckedIntExpr $ bool 222222 111111 early)
+      | otherwise
+      = c v e

testsuite/tests/plugins/simple-plugin/Simple/ReplacePlugin.hs

@@ -51,5 +51,6 @@ fixGuts rep guts = pure $ guts { mg_binds = fmap fix_bind (mg_binds guts) }
       Tick t e      -> Tick t (fix_expr e)
       Type t        -> Type t
       Coercion c    -> Coercion c
+      Let b body    -> Let (fix_bind b) (fix_expr body)
 
     fix_alt (Alt c bs e) = Alt c bs (fix_expr e)