Support TemplateHaskell (#222)

* First attempt at TH support

* Update TcModuleResult when generating core

* Be a bit more cautious when asking for bytecode

* Check need for bytecode not only in source file itself, also in global information

* Add a test (based on #212)

* Fix test (thanks, @jinwoo)

* Split GenerateCore and GenerateByteCode

Author: serras

src/Development/IDE/Core/Compile.hs

@@ -15,6 +15,7 @@ module Development.IDE.Core.Compile
   , computePackageDeps
   , addRelativeImport
   , mkTcModuleResult
+  , generateByteCode
   ) where
 
 import Development.IDE.Core.RuleTypes
@@ -41,6 +42,7 @@ import qualified GHC
 import           GhcMonad
 import           GhcPlugins                     as GHC hiding (fst3, (<>))
 import qualified HeaderInfo                     as Hdr
+import           HscMain                        (hscInteractive)
 import           MkIface
 import           StringBuffer                   as SB
 import           TidyPgm
@@ -122,7 +124,7 @@ compileModule
     :: HscEnv
     -> [TcModuleResult]
     -> TcModuleResult
-    -> IO ([FileDiagnostic], Maybe CoreModule)
+    -> IO ([FileDiagnostic], Maybe (SafeHaskellMode, CgGuts, ModDetails))
 compileModule packageState deps tmr =
     fmap (either (, Nothing) (second Just)) $
     runGhcEnv packageState $
@@ -138,15 +140,22 @@ compileModule packageState deps tmr =
                 GHC.dm_core_module <$> GHC.desugarModule tm'
 
             -- give variables unique OccNames
-            (tidy, details) <- liftIO $ tidyProgram session desugar
+            (guts, details) <- liftIO $ tidyProgram session desugar
+            return (map snd warnings, (mg_safe_haskell desugar, guts, details))
 
-            let core = CoreModule
-                         (cg_module tidy)
-                         (md_types details)
-                         (cg_binds tidy)
-                         (mg_safe_haskell desugar)
-
-            return (map snd warnings, core)
+generateByteCode :: HscEnv -> [TcModuleResult] -> TcModuleResult -> CgGuts -> IO ([FileDiagnostic], Maybe Linkable)
+generateByteCode hscEnv deps tmr guts =
+    fmap (either (, Nothing) (second Just)) $
+    runGhcEnv hscEnv $
+      catchSrcErrors "bytecode" $ do
+          setupEnv (deps ++ [tmr])
+          session <- getSession
+          (warnings, (_, bytecode, sptEntries)) <- withWarnings "bytecode" $ \tweak ->
+              liftIO $ hscInteractive session guts (tweak $ GHC.pm_mod_summary $ GHC.tm_parsed_module $ tmrModule tmr)
+          let summary = pm_mod_summary $ tm_parsed_module $ tmrModule tmr
+          let unlinked = BCOs bytecode sptEntries
+          let linkable = LM (ms_hs_date summary) (ms_mod summary) [unlinked]
+          pure (map snd warnings, linkable)
 
 demoteTypeErrorsToWarnings :: ParsedModule -> ParsedModule
 demoteTypeErrorsToWarnings =

src/Development/IDE/Core/RuleTypes.hs

@@ -24,7 +24,7 @@ import           GHC.Generics                             (Generic)
 
 import           GHC
 import Module (InstalledUnitId)
-import HscTypes (HomeModInfo)
+import HscTypes (CgGuts, Linkable, HomeModInfo, ModDetails)
 import Development.IDE.GHC.Compat
 
 import           Development.IDE.Spans.Type
@@ -65,7 +65,10 @@ type instance RuleResult TypeCheck = TcModuleResult
 type instance RuleResult GetSpanInfo = [SpanInfo]
 
 -- | Convert to Core, requires TypeCheck*
-type instance RuleResult GenerateCore = CoreModule
+type instance RuleResult GenerateCore = (SafeHaskellMode, CgGuts, ModDetails)
+
+-- | Generate byte code for template haskell.
+type instance RuleResult GenerateByteCode = Linkable
 
 -- | A GHC session that we reuse.
 type instance RuleResult GhcSession = HscEnvEq
@@ -131,6 +134,12 @@ instance Hashable GenerateCore
 instance NFData   GenerateCore
 instance Binary   GenerateCore
 
+data GenerateByteCode = GenerateByteCode
+   deriving (Eq, Show, Typeable, Generic)
+instance Hashable GenerateByteCode
+instance NFData   GenerateByteCode
+instance Binary   GenerateByteCode
+
 data GhcSession = GhcSession
     deriving (Eq, Show, Typeable, Generic)
 instance Hashable GhcSession

src/Development/IDE/Core/Rules.hs

@@ -17,7 +17,6 @@ module Development.IDE.Core.Rules(
     runAction, useE, useNoFileE, usesE,
     toIdeResult, defineNoFile,
     mainRule,
-    getGhcCore,
     getAtPoint,
     getDefinition,
     getDependencies,
@@ -55,10 +54,12 @@ import           Development.Shake                        hiding (Diagnostic)
 import Development.IDE.Core.RuleTypes
 
 import           GHC hiding (parseModule, typecheckModule)
+import qualified GHC.LanguageExtensions as LangExt
 import Development.IDE.GHC.Compat
 import           UniqSupply
 import NameCache
 import HscTypes
+import DynFlags (xopt)
 import GHC.Generics(Generic)
 
 import qualified Development.IDE.Spans.AtPoint as AtPoint
@@ -92,16 +93,6 @@ defineNoFile f = define $ \k file -> do
 ------------------------------------------------------------
 -- Exposed API
 
-
--- | Generate the GHC Core for the supplied file and its dependencies.
-getGhcCore :: NormalizedFilePath -> Action (Maybe [CoreModule])
-getGhcCore file = runMaybeT $ do
-    files <- transitiveModuleDeps <$> useE GetDependencies file
-    pms   <- usesE GetParsedModule $ files ++ [file]
-    usesE GenerateCore $ map fileFromParsedModule pms
-
-
-
 -- | Get all transitive file dependencies of a given module.
 -- Does not include the file itself.
 getDependencies :: NormalizedFilePath -> Action (Maybe [NormalizedFilePath])
@@ -281,13 +272,27 @@ typeCheckRule =
     define $ \TypeCheck file -> do
         pm <- use_ GetParsedModule file
         deps <- use_ GetDependencies file
-        tms <- uses_ TypeCheck (transitiveModuleDeps deps)
-        setPriority priorityTypeCheck
         packageState <- hscEnv <$> use_ GhcSession file
+        -- Figure out whether we need TemplateHaskell or QuasiQuotes support
+        let graph_needs_th_qq = needsTemplateHaskellOrQQ $ hsc_mod_graph packageState
+            file_uses_th_qq = uses_th_qq $ ms_hspp_opts (pm_mod_summary pm)
+            any_uses_th_qq = graph_needs_th_qq || file_uses_th_qq
+        tms <- if any_uses_th_qq
+                  -- If we use TH or QQ, we must obtain the bytecode
+                  then do
+                    bytecodes <- uses_ GenerateByteCode (transitiveModuleDeps deps)
+                    tmrs <- uses_ TypeCheck (transitiveModuleDeps deps)
+                    pure (zipWith addByteCode bytecodes tmrs)
+                  else uses_ TypeCheck (transitiveModuleDeps deps)
+        setPriority priorityTypeCheck
         IdeOptions{ optDefer = defer} <- getIdeOptions
         liftIO $ typecheckModule defer packageState tms pm
+    where
+        uses_th_qq dflags = xopt LangExt.TemplateHaskell dflags || xopt LangExt.QuasiQuotes dflags
+        addByteCode :: Linkable -> TcModuleResult -> TcModuleResult
+        addByteCode lm tmr = tmr { tmrModInfo = (tmrModInfo tmr) { hm_linkable = Just lm } }
 
-generateCore :: NormalizedFilePath -> Action (IdeResult CoreModule)
+generateCore :: NormalizedFilePath -> Action (IdeResult (SafeHaskellMode, CgGuts, ModDetails))
 generateCore file = do
     deps <- use_ GetDependencies file
     (tm:tms) <- uses_ TypeCheck (file:transitiveModuleDeps deps)
@@ -299,6 +304,14 @@ generateCoreRule :: Rules ()
 generateCoreRule =
     define $ \GenerateCore -> generateCore
 
+generateByteCodeRule :: Rules ()
+generateByteCodeRule =
+    define $ \GenerateByteCode file -> do
+      deps <- use_ GetDependencies file
+      (tm : tms) <- uses_ TypeCheck (file: transitiveModuleDeps deps)
+      session <- hscEnv <$> use_ GhcSession file
+      (_, guts, _) <- use_ GenerateCore file
+      liftIO $ generateByteCode session tms tm guts
 
 -- A local rule type to get caching. We want to use newCache, but it has
 -- thread killed exception issues, so we lift it to a full rule.
@@ -345,6 +358,7 @@ mainRule = do
     typeCheckRule
     getSpanInfoRule
     generateCoreRule
+    generateByteCodeRule
     loadGhcSession
     getHieFileRule
 

src/Development/IDE/GHC/Orphans.hs

@@ -20,7 +20,13 @@ import Development.IDE.GHC.Util
 -- Orphan instances for types from the GHC API.
 instance Show CoreModule where show = prettyPrint
 instance NFData CoreModule where rnf = rwhnf
-
+instance Show CgGuts where show = prettyPrint . cg_module
+instance NFData CgGuts where rnf = rwhnf
+instance Show ModDetails where show = const "<moddetails>"
+instance NFData ModDetails where rnf = rwhnf
+instance NFData SafeHaskellMode where rnf = rwhnf
+instance Show Linkable where show = prettyPrint
+instance NFData Linkable where rnf = rwhnf
 
 instance Show InstalledUnitId where
     show = installedUnitIdString

src/Development/IDE/GHC/Util.hs

@@ -19,7 +19,8 @@ module Development.IDE.GHC.Util(
     textToStringBuffer,
     moduleImportPath,
     HscEnvEq, hscEnv, newHscEnvEq,
-    readFileUtf8
+    readFileUtf8,
+    cgGutsToCoreModule
     ) where
 
 import Config
@@ -146,3 +147,10 @@ instance NFData HscEnvEq where
 
 readFileUtf8 :: FilePath -> IO T.Text
 readFileUtf8 f = T.decodeUtf8With T.lenientDecode <$> BS.readFile f
+
+cgGutsToCoreModule :: SafeHaskellMode -> CgGuts -> ModDetails -> CoreModule
+cgGutsToCoreModule safeMode guts modDetails = CoreModule
+    (cg_module guts)
+    (md_types modDetails)
+    (cg_binds guts)
+    safeMode

test/exe/Main.hs

@@ -41,6 +41,7 @@ main = defaultMain $ testGroup "HIE"
   , codeActionTests
   , findDefinitionAndHoverTests
   , pluginTests
+  , thTests
   ]
 
 initializeResponseTests :: TestTree
@@ -833,6 +834,36 @@ pluginTests = testSessionWait "plugins" $ do
       )
     ]
 
+thTests :: TestTree
+thTests =
+  testGroup
+    "TemplateHaskell"
+    [ -- Test for https://github.com/digital-asset/ghcide/pull/212
+      testSessionWait "load" $ do
+        let sourceA =
+              T.unlines
+                [ "{-# LANGUAGE PackageImports #-}",
+                  "{-# LANGUAGE TemplateHaskell #-}",
+                  "module A where",
+                  "import \"template-haskell\" Language.Haskell.TH",
+                  "a :: Integer",
+                  "a = $(litE $ IntegerL 3)"
+                ]
+            sourceB =
+              T.unlines
+                [ "{-# LANGUAGE PackageImports #-}",
+                  "{-# LANGUAGE TemplateHaskell #-}",
+                  "module B where",
+                  "import A",
+                  "import \"template-haskell\" Language.Haskell.TH",
+                  "b :: Integer",
+                  "b = $(litE $ IntegerL $ a) + n"
+                ]
+        _ <- openDoc' "A.hs" "haskell" sourceA
+        _ <- openDoc' "B.hs" "haskell" sourceB
+        expectDiagnostics [ ( "B.hs", [(DsError, (6, 29), "Variable not in scope: n")] ) ]
+    ]
+
 xfail :: TestTree -> String -> TestTree
 xfail = flip expectFailBecause