Merge pull request ucsd-progsys#2387 from gergoerdi/ghc-api-cleanup

Improve GHC API usage

liquidhaskell-boot/src-ghc/Liquid/GHC/API.hs

@@ -23,7 +23,7 @@ import           GHC                  as Ghc
     ( Class
     , DataCon
     , DesugaredModule(DesugaredModule, dm_typechecked_module, dm_core_module)
-    , DynFlags(backend, debugLevel, ghcLink, ghcMode)
+    , DynFlags(backend, debugLevel, ghcLink, ghcMode, warningFlags)
     , FixityDirection(InfixN, InfixR)
     , FixitySig(FixitySig)
     , GenLocated(L)
@@ -37,6 +37,7 @@ import           GHC                  as Ghc
     , GhcException(CmdLineError, ProgramError)
     , GhcLink(LinkInMemory)
     , GhcMode(CompManager)
+    , GhcMonad
     , GhcPs
     , GhcRn
     , HsDecl(SigD)
@@ -81,9 +82,11 @@ import           GHC                  as Ghc
     , TypecheckedModule(tm_checked_module_info, tm_internals_, tm_parsed_module)
     , classMethods
     , classSCTheta
+    , coreModule
     , dataConTyCon
     , dataConFieldLabels
     , dataConWrapperType
+    , desugarModule
     , getLocA
     , getLogger
     , getName
@@ -109,6 +112,7 @@ import           GHC                  as Ghc
     , isRecordSelector
     , isTypeSynonymTyCon
     , isVanillaDataCon
+    , lookupName
     , mkHsApp
     , mkHsDictLet
     , mkHsForAllInvisTele
@@ -154,6 +158,7 @@ import           GHC                  as Ghc
     , tyConDataCons
     , tyConKind
     , tyConTyVars
+    , typecheckModule
     , unLoc
     )
 
@@ -412,10 +417,6 @@ import GHC.Driver.Config.Diagnostic as Ghc
     , initDsMessageOpts
     , initIfaceMessageOpts
     )
-import GHC.Driver.Main                as Ghc
-    ( hscDesugar
-    , hscTcRcLookupName
-    )
 import GHC.Driver.Plugins             as Ghc
     ( ParsedResult(..)
     )
@@ -428,7 +429,7 @@ import GHC.Driver.Session             as Ghc
     , updOptLevel
     , xopt_set
     )
-import GHC.Driver.Monad               as Ghc (withSession)
+import GHC.Driver.Monad               as Ghc (withSession, reflectGhc, Session(..))
 import GHC.HsToCore.Monad             as Ghc
     ( DsM, initDsTc, initDsWithModGuts, newUnique )
 import GHC.Iface.Syntax               as Ghc
@@ -452,6 +453,7 @@ import GHC.Driver.Backend             as Ghc (interpreterBackend)
 import GHC.Driver.Env                 as Ghc
     ( HscEnv(hsc_mod_graph, hsc_unit_env, hsc_dflags, hsc_plugins)
     , Hsc
+    , hscSetFlags, hscUpdateFlags
     )
 import GHC.Driver.Errors              as Ghc
     ( printMessages )
@@ -499,6 +501,7 @@ import GHC.Tc.Utils.Monad             as Ghc
     ( captureConstraints
     , discardConstraints
     , getEnv
+    , getTopEnv
     , failIfErrsM
     , failM
     , failWithTc
@@ -510,6 +513,7 @@ import GHC.Tc.Utils.Monad             as Ghc
     , reportDiagnostic
     , reportDiagnostics
     , updEnv
+    , updTopEnv
     )
 import GHC.Tc.Utils.TcType            as Ghc (tcSplitDFunTy, tcSplitMethodTy)
 import GHC.Tc.Zonk.Type               as Ghc

liquidhaskell-boot/src-ghc/Liquid/GHC/API/Extra.hs

@@ -10,14 +10,12 @@ module Liquid.GHC.API.Extra (
   , apiComments
   , apiCommentsParsedSource
   , dataConSig
-  , desugarModuleIO
   , fsToUnitId
   , isPatErrorAlt
   , lookupModSummary
   , minus_RDR
-  , modInfoLookupNameIO
+  , modInfoLookupName
   , moduleInfoTc
-  , parseModuleIO
   , qualifiedNameFS
   , relevantModules
   , renderWithStyle
@@ -28,13 +26,12 @@ module Liquid.GHC.API.Extra (
   , strictNothing
   , thisPackage
   , tyConRealArity
-  , typecheckModuleIO
   , untick
   ) where
 
 import Control.Monad.IO.Class
 import           Liquid.GHC.API.StableModule      as StableModule
-import GHC
+import GHC hiding (modInfoLookupName)
 import Data.Data (Data, gmapQr, gmapT)
 import Data.Generics (extQ, extT)
 import Data.Foldable                  (asum)
@@ -49,7 +46,6 @@ import GHC.Core.Make                  (pAT_ERROR_ID)
 import GHC.Core.Type                  as Ghc hiding (typeKind , isPredTy, extendCvSubst, linear)
 import GHC.Data.Bag                   (bagToList)
 import GHC.Data.FastString            as Ghc
-import qualified GHC.Data.EnumSet as EnumSet
 import GHC.Data.Maybe
 import qualified GHC.Data.Strict
 import GHC.Driver.Env
@@ -63,7 +59,6 @@ import GHC.Types.SourceText           (SourceText(..))
 import GHC.Types.SrcLoc               as Ghc
 import GHC.Types.TypeEnv
 import GHC.Types.Unique               (getUnique, hasKey)
-import GHC.Types.Unique.FM
 
 import GHC.Unit.Module.Deps           as Ghc (Dependencies(dep_direct_mods))
 import GHC.Unit.Module.Graph          as Ghc
@@ -146,52 +141,6 @@ relevantModules mg modGuts = used `S.union` dependencies
           UsageMergedRequirement { usg_mod      = modl    } -> modl : acc
           _ -> acc
 
---
--- Parsing, typechecking and desugaring a module
---
-parseModuleIO :: HscEnv -> ModSummary -> IO ParsedModule
-parseModuleIO hscEnv ms = do
-  let hsc_env_tmp = hscEnv { hsc_dflags = ms_hspp_opts ms }
-  hpm <- hscParse hsc_env_tmp ms
-  return (ParsedModule ms (hpm_module hpm) (hpm_src_files hpm))
-
--- | Our own simplified version of 'TypecheckedModule'.
-data TypecheckedModuleLH = TypecheckedModuleLH {
-    tmlh_parsed_module  :: ParsedModule
-  , tmlh_renamed_source :: Maybe RenamedSource
-  , tmlh_mod_summary    :: ModSummary
-  , tmlh_gbl_env        :: TcGblEnv
-  }
-
-typecheckModuleIO :: HscEnv -> ParsedModule -> IO TypecheckedModuleLH
-typecheckModuleIO hscEnv pmod = do
-  -- Suppress all the warnings, so that they won't be printed (which would result in them being
-  -- printed twice, one by GHC and once here).
-  let ms = pm_mod_summary pmod
-  let dynFlags' = ms_hspp_opts ms
-  let hsc_env_tmp = hscEnv { hsc_dflags = dynFlags' { warningFlags = EnumSet.empty } }
-  (tc_gbl_env, rn_info)
-        <- hscTypecheckRename hsc_env_tmp ms $
-                       HsParsedModule { hpm_module = parsedSource pmod,
-                                        hpm_src_files = pm_extra_src_files pmod }
-  return TypecheckedModuleLH {
-      tmlh_parsed_module  = pmod
-    , tmlh_renamed_source = rn_info
-    , tmlh_mod_summary    = ms
-    , tmlh_gbl_env        = tc_gbl_env
-    }
-
--- | Desugar a typechecked module.
-desugarModuleIO :: HscEnv -> ModSummary -> TypecheckedModuleLH -> IO ModGuts
-desugarModuleIO hscEnv originalModSum typechecked = do
-  -- See [NOTE:ghc810] on why we override the dynFlags here before calling 'desugarModule'.
-  let modSum         = originalModSum { ms_hspp_opts = hsc_dflags hscEnv }
-  let parsedMod'     = (tmlh_parsed_module typechecked) { pm_mod_summary = modSum }
-  let typechecked'   = typechecked { tmlh_parsed_module = parsedMod' }
-
-  let hsc_env_tmp = hscEnv { hsc_dflags = ms_hspp_opts (tmlh_mod_summary typechecked') }
-  hscDesugar hsc_env_tmp (tmlh_mod_summary typechecked') (tmlh_gbl_env typechecked')
-
 -- | Abstraction of 'EpaComment'.
 data ApiComment
   = ApiLineComment String
@@ -275,16 +224,13 @@ lookupModSummary hscEnv mdl = do
 -- | Our own simplified version of 'ModuleInfo' to overcome the fact we cannot construct the \"original\"
 -- one as the constructor is not exported, and 'getHomeModuleInfo' and 'getPackageModuleInfo' are not
 -- exported either, so we had to backport them as well.
-newtype ModuleInfoLH = ModuleInfoLH { minflh_type_env :: UniqFM Name TyThing }
+newtype ModuleInfoLH = ModuleInfoLH { minflh_type_env :: TypeEnv }
 
-modInfoLookupNameIO :: HscEnv
-                  -> ModuleInfoLH
-                  -> Name
-                  -> IO (Maybe TyThing)
-modInfoLookupNameIO hscEnv minf name =
+modInfoLookupName :: (GhcMonad m) => ModuleInfoLH -> Name -> m (Maybe TyThing)
+modInfoLookupName minf name = do
   case lookupTypeEnv (minflh_type_env minf) name of
     Just tyThing -> return (Just tyThing)
-    Nothing      -> lookupType hscEnv name
+    Nothing      -> lookupGlobalName name
 
 moduleInfoTc :: HscEnv -> TcGblEnv -> IO ModuleInfoLH
 moduleInfoTc hscEnv tcGblEnv = do

liquidhaskell-boot/src/Language/Haskell/Liquid/GHC/Interface.hs

@@ -40,8 +40,7 @@ module Language.Haskell.Liquid.GHC.Interface (
   , keepRawTokenStream
   , ignoreInline
   , lookupTyThings
-  , availableTyCons
-  , availableVars
+  , availableTyThings
   , updLiftedSpec
   ) where
 
@@ -184,47 +183,29 @@ qImports qns  = QImports
 --   for this module; we will use this to create our name-resolution environment
 --   (see `Bare.Resolve`)
 ---------------------------------------------------------------------------------------
-lookupTyThings :: HscEnv -> TcGblEnv -> IO [(Name, Maybe TyThing)]
-lookupTyThings hscEnv tcGblEnv = forM names (lookupTyThing hscEnv tcGblEnv)
+lookupTyThings :: (GhcMonad m) => TcGblEnv -> m [(Name, Maybe TyThing)]
+lookupTyThings tcGblEnv = mapM (lookupTyThing tcGblEnv) names
   where
-    names :: [Ghc.Name]
-    names  = liftM2 (++)
-             (fmap Ghc.greName . Ghc.globalRdrEnvElts . tcg_rdr_env)
-             (fmap is_dfun_name . tcg_insts) tcGblEnv
--- | Lookup a single 'Name' in the GHC environment, yielding back the 'Name' alongside the 'TyThing',
--- if one is found.
-lookupTyThing :: HscEnv -> TcGblEnv -> Name -> IO (Name, Maybe TyThing)
-lookupTyThing hscEnv tcGblEnv n = do
-  mty <- runMaybeT $
-         MaybeT (Ghc.hscTcRcLookupName hscEnv n)
-         `mplus`
-         MaybeT (
-           do mi  <- moduleInfoTc hscEnv tcGblEnv
-              modInfoLookupNameIO hscEnv mi n
-           )
-  return (n, mty)
-
-availableTyThings :: HscEnv -> TcGblEnv -> [AvailInfo] -> IO [TyThing]
-availableTyThings hscEnv tcGblEnv avails =
+    names = liftA2 (++)
+        (fmap Ghc.greName . Ghc.globalRdrEnvElts . tcg_rdr_env)
+        (fmap is_dfun_name . tcg_insts)
+        tcGblEnv
+
+lookupTyThing :: (GhcMonad m) => TcGblEnv -> Name -> m (Name, Maybe TyThing)
+lookupTyThing tcGblEnv name = do
+    hscEnv <- getSession
+    mbTy <- runMaybeT . msum . map MaybeT $
+        [ lookupName name
+        , do minf <- liftIO $ moduleInfoTc hscEnv tcGblEnv
+             modInfoLookupName minf name
+        ]
+    return (name, mbTy)
+
+availableTyThings :: (GhcMonad m) => TcGblEnv -> [AvailInfo] -> m [TyThing]
+availableTyThings tcGblEnv avails =
     fmap catMaybes $
-      mapM (fmap snd . lookupTyThing hscEnv tcGblEnv) $
-      availableNames avails
-
--- | Returns all the available (i.e. exported) 'TyCon's (type constructors) for the input 'Module'.
-availableTyCons :: HscEnv -> TcGblEnv -> [AvailInfo] -> IO [Ghc.TyCon]
-availableTyCons hscEnv tcGblEnv avails =
-  fmap (\things -> [tyCon | (ATyCon tyCon) <- things]) (availableTyThings hscEnv tcGblEnv avails)
-
--- | Returns all the available (i.e. exported) 'Var's for the input 'Module'.
-availableVars :: HscEnv -> TcGblEnv -> [AvailInfo] -> IO [Ghc.Var]
-availableVars hscEnv tcGblEnv avails =
-  fmap (\things -> [var | (AnId var) <- things]) (availableTyThings hscEnv tcGblEnv avails)
-
-availableNames :: [AvailInfo] -> [Name]
-availableNames =
-    concatMap $ \case
-      Avail n -> [n]
-      AvailTC n ns -> n : ns
+      mapM (fmap snd . lookupTyThing tcGblEnv) $
+      concatMap availNames avails
 
 _dumpTypeEnv :: TypecheckedModule -> IO ()
 _dumpTypeEnv tm = do

liquidhaskell-boot/src/Language/Haskell/Liquid/GHC/Misc.hs

@@ -623,16 +623,6 @@ stripParens t = fromMaybe t (strip t)
 stripParensSym :: Symbol -> Symbol
 stripParensSym (symbolText -> t) = symbol (stripParens t)
 
-desugarModule :: TypecheckedModule -> Ghc DesugaredModule
-desugarModule tcm = do
-  let ms = pm_mod_summary $ tm_parsed_module tcm
-  -- let ms = modSummary tcm
-  let (tcg, _) = tm_internals_ tcm
-  hsc_env <- getSession
-  let hsc_env_tmp = hsc_env { hsc_dflags = ms_hspp_opts ms }
-  guts <- liftIO $ hscDesugar{- WithLoc -} hsc_env_tmp ms tcg
-  return DesugaredModule { dm_typechecked_module = tcm, dm_core_module = guts }
-
 --------------------------------------------------------------------------------
 -- | GHC Compatibility Layer ---------------------------------------------------
 --------------------------------------------------------------------------------

liquidhaskell-boot/src/Language/Haskell/Liquid/GHC/Plugin.hs

@@ -137,7 +137,7 @@ plugin = GHC.defaultPlugin {
             liftIO $ printWarning logger warning
             pure gblEnv
           else do
-            newGblEnv <- typecheckHook cfg summary gblEnv
+            newGblEnv <- typecheckHook cfg gblEnv
             case newGblEnv of
               -- Exit with success if all expected errors were found
               Left (ErrorsOccurred []) -> pure gblEnv
@@ -279,14 +279,14 @@ parsedHook _ ms parsedResult = do
 --    grab from parsing (again) the module by using the GHC API, so we are really
 --    independent from the \"normal\" compilation pipeline.
 --
-typecheckHook  :: Config -> ModSummary -> TcGblEnv -> TcM (Either LiquidCheckException TcGblEnv)
-typecheckHook cfg0 modSummary0 tcGblEnv = bracket startTypechecking endTypechecking $ \case
+typecheckHook  :: Config -> TcGblEnv -> TcM (Either LiquidCheckException TcGblEnv)
+typecheckHook cfg0 tcGblEnv = bracket startTypechecking endTypechecking $ \case
   Just Typechecking ->
     -- We're being called from the `typecheckModuleIO` call in `typecheckHook`, so we avoid looping
     -- See 'Breadcrumb' for more information.
     pure $ Right tcGblEnv
   Just (Parsed parsed0) ->
-    typecheckHook' cfg0 modSummary0 parsed0 tcGblEnv
+    typecheckHook' cfg0 parsed0 tcGblEnv
   Nothing ->
     -- The module has been verified by an earlier call to the plugin.
     -- This could happen if multiple @-fplugin=LiquidHaskell@ flags are passed to GHC.
@@ -303,14 +303,12 @@ typecheckHook cfg0 modSummary0 tcGblEnv = bracket startTypechecking endTypecheck
         Just Parsed{} -> void $ swapBreadcrumb thisModule Nothing
         _ -> pure ()
 
-typecheckHook' :: Config -> ModSummary -> ParsedModule -> TcGblEnv -> TcM (Either LiquidCheckException TcGblEnv)
-typecheckHook' cfg0 modSummary0 parsed0 tcGblEnv = do
+typecheckHook' :: Config -> ParsedModule -> TcGblEnv -> TcM (Either LiquidCheckException TcGblEnv)
+typecheckHook' cfg0 parsed0 tcGblEnv = do
   debugLog $ "We are in module: " <> show (toStableModule thisModule)
   let modSummary = updateModSummaryDynFlags unoptimiseDynFlags modSummary0
       thisFile = LH.modSummaryHsFile modSummary
 
-  env0 <- env_top <$> getEnv
-  let env = env0 { hsc_dflags = ms_hspp_opts modSummary }
   let specComments = map mkSpecComment $ LH.extractSpecComments parsed0
       parsed = addNoInlinePragmasToLocalBinds parsed0
 
@@ -323,29 +321,38 @@ typecheckHook' cfg0 modSummary0 parsed0 tcGblEnv = do
 
         let modSummary2 = updateModSummaryDynFlags (maybeInsertBreakPoints cfg) modSummary
             parsed2 = parsed { pm_mod_summary = modSummary2 }
-            env2 = env { hsc_dflags = ms_hspp_opts modSummary2 }
 
-        typechecked     <- liftIO $ typecheckModuleIO env2 (LH.ignoreInline parsed2)
-        resolvedNames   <- liftIO $ LH.lookupTyThings env2 tcGblEnv
-        availTyCons     <- liftIO $ LH.availableTyCons env2 tcGblEnv (tcg_exports tcGblEnv)
-        availVars       <- liftIO $ LH.availableVars env2 tcGblEnv (tcg_exports tcGblEnv)
+        updTopEnv (hscUpdateFlags noWarnings . hscSetFlags (ms_hspp_opts modSummary2)) $ do
+          env2 <- getTopEnv
 
-        unoptimisedGuts <- liftIO $ desugarModuleIO env2 modSummary2 typechecked
+          pipelineData <- liftIO $ do
+              session <- Session <$> newIORef env2
+              flip reflectGhc session $ do
+                  typechecked     <- typecheckModule (LH.ignoreInline parsed2)
+                  unoptimisedGuts <- desugarModule typechecked
 
-        let tcData = mkTcData (tcg_rn_imports tcGblEnv) resolvedNames availTyCons availVars
-        let pipelineData = PipelineData unoptimisedGuts tcData specs
+                  resolvedNames   <- LH.lookupTyThings tcGblEnv
+                  avails          <- LH.availableTyThings tcGblEnv (tcg_exports tcGblEnv)
+                  let availTyCons = [ tc | ATyCon tc <- avails ]
+                      availVars   = [ var | AnId var <- avails ]
+
+                  let tcData = mkTcData (tcg_rn_imports tcGblEnv) resolvedNames availTyCons availVars
+                  return $ PipelineData (coreModule unoptimisedGuts) tcData specs
 
-        updEnv (\e -> e {env_top = env2}) $
           liquidHaskellCheckWithConfig cfg pipelineData modSummary2 tcGblEnv
 
   where
     thisModule :: Module
     thisModule = tcg_mod tcGblEnv
 
+    modSummary0 = pm_mod_summary parsed0
+
     continue = pure $ Left (ErrorsOccurred [])
 
     updateModSummaryDynFlags f ms = ms { ms_hspp_opts = f (ms_hspp_opts ms) }
 
+    noWarnings dflags = dflags { warningFlags = mempty }
+
 serialiseSpec :: Module -> TcGblEnv -> LiquidLib -> TcM TcGblEnv
 serialiseSpec thisModule tcGblEnv liquidLib = do
   -- ---