Add GetHieAsts rule

Author: wz1000

src-ghc810/Development/IDE/GHC/HieAst.hs

@@ -17,7 +17,7 @@ Main functions for .hie file generation
 {-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE DeriveDataTypeable #-}
-module Development.IDE.GHC.HieAst ( mkHieFile ) where
+module Development.IDE.GHC.HieAst ( mkHieFile, enrichHie ) where
 
 import GhcPrelude
 

src-ghc86/Development/IDE/GHC/HieAst.hs

@@ -17,7 +17,7 @@ Main functions for .hie file generation
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE DataKinds #-}
-module Development.IDE.GHC.HieAst ( mkHieFile ) where
+module Development.IDE.GHC.HieAst ( mkHieFile, enrichHie ) where
 
 import Avail                      ( Avails )
 import Bag                        ( Bag, bagToList )

src-ghc88/Development/IDE/GHC/HieAst.hs

@@ -16,7 +16,7 @@ Main functions for .hie file generation
 {-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE DeriveDataTypeable #-}
-module Development.IDE.GHC.HieAst ( mkHieFile ) where
+module Development.IDE.GHC.HieAst ( mkHieFile, enrichHie ) where
 
 import Avail                      ( Avails )
 import Bag                        ( Bag, bagToList )

src/Development/IDE/Core/Compile.hs

@@ -18,7 +18,8 @@ module Development.IDE.Core.Compile
   , addRelativeImport
   , mkTcModuleResult
   , generateByteCode
-  , generateAndWriteHieFile
+  , generateHieAsts
+  , writeHieFile
   , writeHiFile
   , getModSummaryFromImports
   , loadHieFile
@@ -55,7 +56,7 @@ import ErrUtils
 #endif
 
 import           Finder
-import           Development.IDE.GHC.Compat hiding (parseModule, typecheckModule)
+import           Development.IDE.GHC.Compat hiding (parseModule, typecheckModule, writeHieFile)
 import qualified Development.IDE.GHC.Compat     as GHC
 import qualified Development.IDE.GHC.Compat     as Compat
 import           GhcMonad
@@ -64,7 +65,7 @@ import qualified HeaderInfo                     as Hdr
 import           HscMain                        (hscInteractive, hscSimplify)
 import           MkIface
 import           StringBuffer                   as SB
-import           TcRnMonad (tct_id, TcTyThing(AGlobal, ATcId), initTc, initIfaceLoad, tcg_th_coreplugins)
+import           TcRnMonad (tct_id, TcTyThing(AGlobal, ATcId), initTc, initIfaceLoad, tcg_th_coreplugins, tcg_binds)
 import           TcIface                        (typecheckIface)
 import           TidyPgm
 
@@ -270,7 +271,7 @@ mkTcModuleResult tcm upgradedError = do
     (iface, _) <- liftIO $ mkIfaceTc session Nothing sf details tcGblEnv
 #endif
     let mod_info = HomeModInfo iface details Nothing
-    return $ TcModuleResult tcm mod_info upgradedError
+    return $ TcModuleResult tcm mod_info upgradedError Nothing
   where
     (tcGblEnv, details) = tm_internals_ tcm
 
@@ -281,19 +282,25 @@ atomicFileWrite targetPath write = do
   (tempFilePath, cleanUp) <- newTempFileWithin dir
   (write tempFilePath >> renameFile tempFilePath targetPath) `onException` cleanUp
 
-generateAndWriteHieFile :: HscEnv -> TypecheckedModule -> BS.ByteString -> IO [FileDiagnostic]
-generateAndWriteHieFile hscEnv tcm source =
-  handleGenerationErrors dflags "extended interface generation" $ do
+generateHieAsts :: HscEnv -> TypecheckedModule -> IO ([FileDiagnostic], Maybe (HieASTs Type))
+generateHieAsts hscEnv tcm =
+  handleGenerationErrors' dflags "extended interface generation" $ do
     case tm_renamed_source tcm of
-      Just rnsrc -> do
-        hf <- runHsc hscEnv $
-          GHC.mkHieFile mod_summary (fst $ tm_internals_ tcm) rnsrc source
-        atomicFileWrite targetPath $ flip GHC.writeHieFile hf
+      Just rnsrc -> runHsc hscEnv $
+        Just <$> GHC.enrichHie (tcg_binds $ fst $ tm_internals_ tcm) rnsrc
       _ ->
-        return ()
+        return Nothing
+  where
+    dflags       = hsc_dflags hscEnv
+
+writeHieFile :: HscEnv -> ModSummary -> [GHC.AvailInfo] -> HieASTs Type -> BS.ByteString -> IO [FileDiagnostic]
+writeHieFile hscEnv mod_summary exports ast source =
+  handleGenerationErrors dflags "extended interface write/compression" $ do
+    hf <- runHsc hscEnv $
+      GHC.mkHieFile' mod_summary exports ast source
+    atomicFileWrite targetPath $ flip GHC.writeHieFile hf
   where
     dflags       = hsc_dflags hscEnv
-    mod_summary  = pm_mod_summary $ tm_parsed_module tcm
     mod_location = ms_location mod_summary
     targetPath   = Compat.ml_hie_file mod_location
 
@@ -315,6 +322,14 @@ handleGenerationErrors dflags source action =
     . (("Error during " ++ T.unpack source) ++) . show @SomeException
     ]
 
+handleGenerationErrors' :: DynFlags -> T.Text -> IO (Maybe a) -> IO ([FileDiagnostic], Maybe a)
+handleGenerationErrors' dflags source action =
+  fmap ([],) action `catches`
+    [ Handler $ return . (,Nothing) . diagFromGhcException source dflags
+    , Handler $ return . (,Nothing) . diagFromString source DsError (noSpan "<internal>")
+    . (("Error during " ++ T.unpack source) ++) . show @SomeException
+    ]
+
 
 -- | Setup the environment that GHC needs according to our
 -- best understanding (!)

src/Development/IDE/Core/RuleTypes.hs

@@ -65,6 +65,7 @@ data TcModuleResult = TcModuleResult
     -- HomeModInfo instead
     , tmrModInfo    :: HomeModInfo
     , tmrDeferedError :: !Bool -- ^ Did we defer any type errors for this module?
+    , tmrHieAsts :: !(Maybe (HieASTs Type)) -- ^ The HieASTs if we computed them
     }
 instance Show TcModuleResult where
     show = show . pm_mod_summary . tm_parsed_module . tmrModule
@@ -97,9 +98,21 @@ instance NFData HiFileResult where
 instance Show HiFileResult where
     show = show . hirModSummary
 
+-- | Save the uncompressed AST here, we compress it just before writing to disk
+data HieAstResult = HAR { hieModule :: Module, hieAst :: !(HieASTs Type), refMap :: !RefMap }
+ 
+instance NFData HieAstResult where
+    rnf (HAR m hf rm) = rnf m `seq` rwhnf hf `seq` rnf rm
+ 
+instance Show HieAstResult where
+    show = show . hieModule
+
 -- | The type checked version of this file, requires TypeCheck+
 type instance RuleResult TypeCheck = TcModuleResult
 
+-- | The uncompressed HieAST
+type instance RuleResult GetHieAst = HieAstResult
+ 
 -- | Information about what spans occur where, requires TypeCheck
 type instance RuleResult GetSpanInfo = SpansInfo
 
@@ -201,6 +214,12 @@ instance Hashable GetSpanInfo
 instance NFData   GetSpanInfo
 instance Binary   GetSpanInfo
 
+data GetHieAst = GetHieAst
+    deriving (Eq, Show, Typeable, Generic)
+instance Hashable GetHieAst
+instance NFData   GetHieAst
+instance Binary   GetHieAst
+
 data GenerateCore = GenerateCore
     deriving (Eq, Show, Typeable, Generic)
 instance Hashable GenerateCore

src/Development/IDE/Core/Rules.hs

@@ -46,7 +46,7 @@ import           Development.IDE.Core.FileExists
 import           Development.IDE.Core.FileStore        (modificationTime, getFileContents)
 import           Development.IDE.Types.Diagnostics as Diag
 import Development.IDE.Types.Location
-import Development.IDE.GHC.Compat hiding (parseModule, typecheckModule)
+import Development.IDE.GHC.Compat hiding (parseModule, typecheckModule, writeHieFile)
 import Development.IDE.GHC.Util
 import Development.IDE.GHC.WithDynFlags
 import Data.Either.Extra
@@ -504,6 +504,21 @@ getDependenciesRule =
         let mbFingerprints = map (fingerprintString . fromNormalizedFilePath) allFiles <$ optShakeFiles opts
         return (fingerprintToBS . fingerprintFingerprints <$> mbFingerprints, ([], transitiveDeps depInfo file))
 
+getHieAstsRule :: Rules ()
+getHieAstsRule =
+    define $ \GetHieAst f -> do
+      tmr <- use_ TypeCheck f
+      (diags,masts) <- case tmrHieAsts tmr of
+        -- If we already have them from typechecking, return them
+        Just asts -> pure ([], Just asts)
+        -- Compute asts if we haven't already computed them
+        Nothing -> do
+          hsc <- hscEnv <$> use_ GhcSession f
+          (diagsHieGen, masts) <- liftIO $ generateHieAsts hsc (tmrModule tmr)
+          pure (diagsHieGen, masts)
+      let refmap = generateReferencesMap . getAsts <$> masts
+      pure (diags, HAR (ms_mod  $ tmrModSummary tmr) <$> masts <*> refmap)
+
 -- Source SpanInfo is used by AtPoint and Goto Definition.
 getSpanInfoRule :: Rules ()
 getSpanInfoRule =
@@ -570,14 +585,20 @@ typeCheckRuleDefinition hsc pm isFoi source = do
         case isFoi of
           IsFOI Modified -> return (diags, Just tcm)
           _ -> do -- If the file is saved on disk, or is not a FOI, we write out ifaces
-            diagsHie <- generateAndWriteHieFile hsc (tmrModule tcm) (fromMaybe "" source)
+            let tm = tmrModule tcm
+                ms = tmrModSummary tcm
+                exports = tcg_exports $ fst $ tm_internals_ tm
+            (diagsHieGen, masts) <- generateHieAsts hsc (tmrModule tcm)
+            diagsHieWrite <- case masts of
+              Nothing -> pure mempty
+              Just asts -> writeHieFile hsc ms exports asts $ fromMaybe "" source
             -- Don't save interface files for modules that compiled due to defering
             -- type errors, as we won't get proper diagnostics if we load these from
             -- disk
             diagsHi  <- if not $ tmrDeferedError tcm
                         then writeHiFile hsc tcm
                         else pure mempty
-            return (diags <> diagsHi <> diagsHie, Just tcm)
+            return (diags <> diagsHi <> diagsHieGen <> diagsHieWrite, Just tcm{tmrHieAsts = masts})
       (diags, res) ->
         return (diags, snd <$> res)
  where
@@ -849,6 +870,7 @@ mainRule = do
     isHiFileStableRule
     getModuleGraphRule
     knownFilesRule
+    getHieAstsRule
 
 -- | Given the path to a module src file, this rule returns True if the
 -- corresponding `.hi` file is stable, that is, if it is newer

src/Development/IDE/GHC/Compat.hs

@@ -16,6 +16,9 @@ module Development.IDE.GHC.Compat(
     NameCacheUpdater(..),
     hieExportNames,
     mkHieFile,
+    mkHieFile',
+    enrichHie,
+    RefMap,
     writeHieFile,
     readHieFile,
     supportsHieFiles,
@@ -54,6 +57,8 @@ module Development.IDE.GHC.Compat(
     getLoc,
     upNameCache,
     disableWarningsAsErrors,
+    AvailInfo,
+    tcg_exports,
 
     module GHC,
 #if MIN_GHC_API_VERSION(8,6,0)
@@ -78,6 +83,10 @@ import Packages
 import Data.IORef
 import HscTypes
 import NameCache
+import qualified Data.ByteString as BS
+import MkIface
+import Data.Map.Strict (Map)
+import TcRnTypes
 
 import qualified GHC
 import GHC hiding (
@@ -110,7 +119,7 @@ import ErrUtils (ErrorMessages)
 import FastString (FastString)
 
 #if MIN_GHC_API_VERSION(8,6,0)
-import Development.IDE.GHC.HieAst (mkHieFile)
+import Development.IDE.GHC.HieAst (mkHieFile,enrichHie)
 import Development.IDE.GHC.HieBin
 
 #if MIN_GHC_API_VERSION(8,8,0)
@@ -181,6 +190,25 @@ includePathsGlobal = id
 includePathsQuote = const []
 #endif
 
+type RefMap = Map Identifier [(Span, IdentifierDetails Type)]
+
+mkHieFile' :: ModSummary
+           -> [AvailInfo]
+           -> HieASTs Type
+           -> BS.ByteString
+           -> Hsc HieFile
+mkHieFile' ms exports asts src = do
+  let Just src_file = ml_hs_file $ ms_location ms
+      (asts',arr) = compressTypes asts
+  return $ HieFile
+      { hie_hs_file = src_file
+      , hie_module = ms_mod ms
+      , hie_types = arr
+      , hie_asts = asts'
+      -- mkIfaceExports sorts the AvailInfos for stability
+      , hie_exports = mkIfaceExports exports
+      , hie_hs_src = src
+      }
 
 addIncludePathsQuote :: FilePath -> DynFlags -> DynFlags
 #if MIN_GHC_API_VERSION(8,6,0)

src/Development/IDE/GHC/Orphans.hs

@@ -80,3 +80,12 @@ instance Show ModuleName where
     show = moduleNameString
 instance Hashable ModuleName where
     hashWithSalt salt = hashWithSalt salt . show
+
+instance NFData a => NFData (IdentifierDetails a) where
+    rnf (IdentifierDetails a b) = rnf a `seq` rnf (length b)
+ 
+instance NFData RealSrcSpan where
+    rnf = rwhnf
+
+instance NFData Type where
+    rnf = rwhnf