Implement TH addCorePlugin.

This allows template-haskell code to add plugins to the compilation
pipeline. Otherwise, the user would have to pass -fplugin=... to ghc.

For now, plugin modules in the current package can't be used. This is
because when TH runs, it is too late to let GHC know that the plugin
modules needed to be compiled first.

Test Plan: ./validate

Reviewers: simonpj, bgamari, austin, goldfire

Reviewed By: bgamari

Subscribers: angerman, rwbarton, mboes, thomie

GHC Trac Issues: #13608

Differential Revision: https://phabricator.haskell.org/D3821

6 files changed, 46 insertions, 11 deletions

diff --git a/compiler/main/DynFlags.hs b/compiler/main/DynFlags.hs
index 0bc54be7d1..4c62a0d464 100644
--- a/compiler/main/DynFlags.hs
+++ b/compiler/main/DynFlags.hs
@@ -91,6 +91,7 @@ module DynFlags (
         opt_windres, opt_lo, opt_lc, opt_lcc,
 
         -- ** Manipulating DynFlags
+        addPluginModuleName,
         defaultDynFlags,                -- Settings -> DynFlags
         defaultWays,
         interpWays,
diff --git a/compiler/main/GHC.hs b/compiler/main/GHC.hs
index 63c15518b5..343ef37210 100644
--- a/compiler/main/GHC.hs
+++ b/compiler/main/GHC.hs
@@ -1032,16 +1032,19 @@ compileCore simplify fn = do
      Just modSummary -> do
        -- Now we have the module name;
        -- parse, typecheck and desugar the module
-       mod_guts <- coreModule `fmap`
-                      -- TODO: space leaky: call hsc* directly?
-                      (desugarModule =<< typecheckModule =<< parseModule modSummary)
+       (tcg, mod_guts) <- -- TODO: space leaky: call hsc* directly?
+         do tm <- typecheckModule =<< parseModule modSummary
+            let tcg = fst (tm_internals tm)
+            (,) tcg . coreModule <$> desugarModule tm
        liftM (gutsToCoreModule (mg_safe_haskell mod_guts)) $
          if simplify
           then do
              -- If simplify is true: simplify (hscSimplify), then tidy
              -- (tidyProgram).
              hsc_env <- getSession
-             simpl_guts <- liftIO $ hscSimplify hsc_env mod_guts
+             simpl_guts <- liftIO $ do
+               plugins <- readIORef (tcg_th_coreplugins tcg)
+               hscSimplify hsc_env plugins mod_guts
              tidy_guts <- liftIO $ tidyProgram hsc_env simpl_guts
              return $ Left tidy_guts
           else
diff --git a/compiler/main/HscMain.hs b/compiler/main/HscMain.hs
index f7a7933db4..8040b1dfb2 100644
--- a/compiler/main/HscMain.hs
+++ b/compiler/main/HscMain.hs
@@ -85,6 +85,7 @@ module HscMain
 import GhcPrelude
 
 import Data.Data hiding (Fixity, TyCon)
+import DynFlags         (addPluginModuleName)
 import Id
 import GHCi             ( addSptEntry )
 import GHCi.RemoteTypes ( ForeignHValue )
@@ -753,7 +754,8 @@ finish hsc_env summary tc_result mb_old_hash = do
           -- and generate a simple interface.
           then mk_simple_iface
           else do
-            desugared_guts <- hscSimplify' desugared_guts0
+            plugins <- liftIO $ readIORef (tcg_th_coreplugins tc_result)
+            desugared_guts <- hscSimplify' plugins desugared_guts0
             (iface, changed, details, cgguts) <-
               liftIO $ hscNormalIface hsc_env desugared_guts mb_old_hash
             return (iface, changed, details, HscRecomp cgguts summary)
@@ -1188,14 +1190,18 @@ hscGetSafeMode tcg_env = do
 -- Simplifiers
 --------------------------------------------------------------
 
-hscSimplify :: HscEnv -> ModGuts -> IO ModGuts
-hscSimplify hsc_env modguts = runHsc hsc_env $ hscSimplify' modguts
+hscSimplify :: HscEnv -> [String] -> ModGuts -> IO ModGuts
+hscSimplify hsc_env plugins modguts =
+    runHsc hsc_env $ hscSimplify' plugins modguts
 
-hscSimplify' :: ModGuts -> Hsc ModGuts
-hscSimplify' ds_result = do
+hscSimplify' :: [String] -> ModGuts -> Hsc ModGuts
+hscSimplify' plugins ds_result = do
     hsc_env <- getHscEnv
+    let hsc_env_with_plugins = hsc_env
+          { hsc_dflags = foldr addPluginModuleName (hsc_dflags hsc_env) plugins
+          }
     {-# SCC "Core2Core" #-}
-      liftIO $ core2core hsc_env ds_result
+      liftIO $ core2core hsc_env_with_plugins ds_result
 
 --------------------------------------------------------------
 -- Interface generators
@@ -1578,7 +1584,9 @@ hscDeclsWithLocation hsc_env0 str source linenumber =
     ds_result <- hscDesugar' iNTERACTIVELoc tc_gblenv
 
     {- Simplify -}
-    simpl_mg <- liftIO $ hscSimplify hsc_env ds_result
+    simpl_mg <- liftIO $ do
+      plugins <- readIORef (tcg_th_coreplugins tc_gblenv)
+      hscSimplify hsc_env plugins ds_result
 
     {- Tidy -}
     (tidy_cg, mod_details) <- liftIO $ tidyProgram hsc_env simpl_mg
diff --git a/compiler/typecheck/TcRnMonad.hs b/compiler/typecheck/TcRnMonad.hs
index fba0f5d379..0e88e237e6 100644
--- a/compiler/typecheck/TcRnMonad.hs
+++ b/compiler/typecheck/TcRnMonad.hs
@@ -221,6 +221,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
         th_foreign_files_var <- newIORef [] ;
         th_topnames_var      <- newIORef emptyNameSet ;
         th_modfinalizers_var <- newIORef [] ;
+        th_coreplugins_var <- newIORef [] ;
         th_state_var         <- newIORef Map.empty ;
         th_remote_state_var  <- newIORef Nothing ;
         let {
@@ -237,6 +238,7 @@ initTc hsc_env hsc_src keep_rn_syntax mod loc do_this
                 tcg_th_foreign_files = th_foreign_files_var,
                 tcg_th_topnames      = th_topnames_var,
                 tcg_th_modfinalizers = th_modfinalizers_var,
+                tcg_th_coreplugins = th_coreplugins_var,
                 tcg_th_state         = th_state_var,
                 tcg_th_remote_state  = th_remote_state_var,
 
diff --git a/compiler/typecheck/TcRnTypes.hs b/compiler/typecheck/TcRnTypes.hs
index a29ad92f18..4c708dd8a2 100644
--- a/compiler/typecheck/TcRnTypes.hs
+++ b/compiler/typecheck/TcRnTypes.hs
@@ -651,6 +651,9 @@ data TcGblEnv
         -- They are computations in the @TcM@ monad rather than @Q@ because we
         -- set them to use particular local environments.
 
+        tcg_th_coreplugins :: TcRef [String],
+        -- ^ Core plugins added by Template Haskell code.
+
         tcg_th_state :: TcRef (Map TypeRep Dynamic),
         tcg_th_remote_state :: TcRef (Maybe (ForeignRef (IORef QState))),
         -- ^ Template Haskell state
diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index 9e102130af..a01022744c 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -35,6 +35,7 @@ import GhcPrelude
 
 import HsSyn
 import Annotations
+import Finder
 import Name
 import TcRnMonad
 import TcType
@@ -920,6 +921,22 @@ instance TH.Quasi TcM where
       fref <- liftIO $ mkForeignRef r (freeRemoteRef r)
       addModFinalizerRef fref
 
+  qAddCorePlugin plugin = do
+      hsc_env <- env_top <$> getEnv
+      r <- liftIO $ findHomeModule hsc_env (mkModuleName plugin)
+      let err = hang
+            (text "addCorePlugin: invalid plugin module "
+               <+> text (show plugin)
+            )
+            2
+            (text "Plugins in the current package can't be specified.")
+      case r of
+        Found {} -> addErr err
+        FoundMultiple {} -> addErr err
+        _ -> return ()
+      th_coreplugins_var <- tcg_th_coreplugins <$> getGblEnv
+      updTcRef th_coreplugins_var (plugin:)
+
   qGetQ :: forall a. Typeable a => TcM (Maybe a)
   qGetQ = do
       th_state_var <- fmap tcg_th_state getGblEnv
@@ -1104,6 +1121,7 @@ handleTHMessage msg = case msg of
   AddModFinalizer r -> do
     hsc_env <- env_top <$> getEnv
     wrapTHResult $ liftIO (mkFinalizedHValue hsc_env r) >>= addModFinalizerRef
+  AddCorePlugin str -> wrapTHResult $ TH.qAddCorePlugin str
   AddTopDecls decs -> wrapTHResult $ TH.qAddTopDecls decs
   AddForeignFile lang str -> wrapTHResult $ TH.qAddForeignFile lang str
   IsExtEnabled ext -> wrapTHResult $ TH.qIsExtEnabled ext