Optimise module to filename

flake.nix

@@ -60,7 +60,7 @@
             chmod +x $dest
           '';
 
-        mkDevShell = hpkgs: with pkgs; mkShell {
+        mkDevShell = hpkgs: with pkgs; pkgs.mkShell {
           name = "haskell-language-server-dev-ghc${hpkgs.ghc.version}";
           # For binary Haskell tools, we use the default Nixpkgs GHC version.
           # This removes a rebuild with a different GHC version. The drawback of
@@ -106,7 +106,9 @@
       in {
         # Developement shell with only dev tools
         devShells = {
-          default = mkDevShell pkgs.haskellPackages;
+          default = pkgs.mkShell {
+            buildInputs = with pkgs; [zlib haskell.compiler.ghc910 cabal-install];
+          };
           shell-ghc96 = mkDevShell pkgs.haskell.packages.ghc96;
           shell-ghc98 = mkDevShell pkgs.haskell.packages.ghc98;
           shell-ghc910 = mkDevShell pkgs.haskell.packages.ghc910;

ghcide-test/exe/CompletionTests.hs

@@ -32,6 +32,7 @@ import           Test.Hls.FileSystem            (file, text)
 import           Test.Hls.Util
 import           Test.Tasty
 import           Test.Tasty.HUnit
+import Debug.Pretty.Simple
 
 tests :: TestTree
 tests

ghcide/ghcide.cabal

@@ -108,6 +108,7 @@ library
     , unliftio-core
     , unordered-containers         >=0.2.10.0
     , vector
+    , pretty-simple
 
   if os(windows)
     build-depends: Win32

ghcide/session-loader/Development/IDE/Session.hs

@@ -91,7 +91,8 @@
 import           Data.Void
 
 import           Control.Concurrent.STM.Stats        (atomically, modifyTVar',
-                                                      readTVar, writeTVar)
+                                                      readTVar, readTVarIO,
+                                                      writeTVar)
 import           Control.Concurrent.STM.TQueue
 import           Control.DeepSeq
 import           Control.Exception                   (evaluate)
@@ -119,6 +120,7 @@
 import           System.Random                       (RandomGen)
 import           Text.ParserCombinators.ReadP        (readP_to_S)
 
+import           Development.IDE                     (HscEnvEq (..))
 import           GHC.Driver.Env                      (hsc_all_home_unit_ids)
 import           GHC.Driver.Errors.Types
 import           GHC.Types.Error                     (errMsgDiagnostic,
@@ -443,7 +445,7 @@
 
   return $ do
     clientConfig <- getClientConfigAction
-    extras@ShakeExtras{restartShakeSession, ideNc, knownTargetsVar, lspEnv
+    extras@ShakeExtras{restartShakeSession, ideNc, knownTargetsVar, lspEnv, moduleToPathCache
                       } <- getShakeExtras
     let invalidateShakeCache = do
             void $ modifyVar' version succ
@@ -476,8 +478,9 @@
                 -- If we don't generate a TargetFile for each potential location, we will only have
                 -- 'TargetFile Foo.hs' in the 'knownTargetsVar', thus not find 'TargetFile Foo.hs-boot'
                 -- and also not find 'TargetModule Foo'.
-                fs <- filterM (IO.doesFileExist . fromNormalizedFilePath) targetLocations
-                pure $ map (\fp -> (TargetFile fp, Set.singleton fp)) (nubOrd (f:fs))
+                pure $ do
+                  file <- nubOrd (f:targetLocations)
+                  pure $ (TargetFile file, Set.singleton file)
               TargetModule _ -> do
                 found <- filterM (IO.doesFileExist . fromNormalizedFilePath) targetLocations
                 return [(targetTarget, Set.fromList found)]
@@ -489,6 +492,22 @@
             pure hasUpdate
           for_ hasUpdate $ \x ->
             logWith recorder Debug $ LogKnownFilesUpdated (targetMap x)
+
+
+          -- Clean the module map cache
+          -- TODO: the clean is total: it refresh the complete module to
+          -- filename cache. We can imagine something smarter in the future,
+          -- but anyway, the scan is actually really fast (It lists recursively
+          -- the content of all your include path, but once. It could only be
+          -- as slow as the number of files in your include paths, which is,
+          -- most of the time, the same as the number of module in your
+          -- project. If there are a lot of not required files inside your
+          -- include path, this will be an issue) and right now
+          -- what's expensive is the association of Known target to module,
+          -- which is still fast considering that it does not do any IO.
+          atomically $ do
+            writeTVar moduleToPathCache mempty
+
           return $ toNoFileKey GetKnownTargets
 
     -- Create a new HscEnv from a hieYaml root and a set of options
@@ -629,7 +648,7 @@
                  [] -> error $ "GHC version could not be parsed: " <> version
                  ((runTime, _):_)
                    | compileTime == runTime -> do
                      atomicModifyIORef' cradle_files (\xs -> (cfp:xs,()))
                      session (hieYaml, toNormalizedFilePath' cfp, opts, libDir)
                    | otherwise -> return (([renderPackageSetupException cfp GhcVersionMismatch{..}], Nothing),[])
              -- Failure case, either a cradle error or the none cradle
@@ -764,6 +783,7 @@
       targetDepends   :: !DependencyInfo,
       targetLocations :: ![NormalizedFilePath]
   }
+  deriving (Show)
 
 fromTargetId :: [FilePath]          -- ^ import paths
              -> [String]            -- ^ extensions to consider
@@ -896,7 +916,7 @@
             x <- map errMsgDiagnostic closure_errs
             DriverHomePackagesNotClosed us <- pure x
             pure us
         isBad ci = (homeUnitId_ (componentDynFlags ci)) `OS.member` bad_units
     -- Whenever we spin up a session on Linux, dynamically load libm.so.6
     -- in. We need this in case the binary is statically linked, in which
     -- case the interactive session will fail when trying to load

ghcide/src/Development/IDE/Core/RuleTypes.hs

@@ -412,6 +412,9 @@ type instance RuleResult GetModSummary = ModSummaryResult
 -- | Generate a ModSummary with the timestamps and preprocessed content elided, for more successful early cutoff
 type instance RuleResult GetModSummaryWithoutTimestamps = ModSummaryResult
 
+type instance RuleResult GetModulesPaths = (M.Map ModuleName (UnitId, NormalizedFilePath),
+                                            M.Map ModuleName (UnitId, NormalizedFilePath))
+
 data GetParsedModule = GetParsedModule
     deriving (Eq, Show, Generic)
 instance Hashable GetParsedModule
@@ -524,6 +527,13 @@ data GetModSummaryWithoutTimestamps = GetModSummaryWithoutTimestamps
 instance Hashable GetModSummaryWithoutTimestamps
 instance NFData   GetModSummaryWithoutTimestamps
 
+-- | Scan all the import directory for existing modules and build a map from
+-- module name to paths
+data GetModulesPaths = GetModulesPaths
+    deriving (Eq, Show, Generic)
+instance Hashable GetModulesPaths
+instance NFData   GetModulesPaths
+
 data GetModSummary = GetModSummary
     deriving (Eq, Show, Generic)
 instance Hashable GetModSummary

ghcide/src/Development/IDE/Core/Rules.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE CPP                   #-}
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE TypeFamilies          #-}
+{-# LANGUAGE PartialTypeSignatures #-}
 
 -- | A Shake implementation of the compiler service, built
 --   using the "Shaker" abstraction layer for in-memory use.
@@ -55,6 +56,7 @@ module Development.IDE.Core.Rules(
     GhcSessionDepsConfig(..),
     Log(..),
     DisplayTHWarning(..),
+    extendModuleMapWithKnownTargets,
     ) where
 
 import           Control.Applicative
@@ -175,6 +177,10 @@ import           System.Info.Extra                            (isWindows)
 
 import qualified Data.IntMap                                  as IM
 import           GHC.Fingerprint
+import qualified Data.Map.Strict as Map
+import System.FilePath (takeExtension, takeFileName, normalise, dropExtension, splitDirectories, equalFilePath)
+import Data.Char (isUpper)
+import System.Directory.Extra (listFilesInside)
 
 data Log
   = LogShake Shake.Log
@@ -313,30 +319,22 @@ getParsedModuleDefinition packageState opt file ms = do
 getLocatedImportsRule :: Recorder (WithPriority Log) -> Rules ()
 getLocatedImportsRule recorder =
     define (cmapWithPrio LogShake recorder) $ \GetLocatedImports file -> do
+
         ModSummaryResult{msrModSummary = ms} <- use_ GetModSummaryWithoutTimestamps file
-        (KnownTargets targets targetsMap) <- useNoFile_ GetKnownTargets
+        -- TODO: should we reverse this concatenation, there are way less
+        -- source import than normal import in theory, so it should be faster
         let imports = [(False, imp) | imp <- ms_textual_imps ms] ++ [(True, imp) | imp <- ms_srcimps ms]
         env_eq <- use_ GhcSession file
         let env = hscEnv env_eq
         let import_dirs = map (second homeUnitEnv_dflags) $ hugElts $ hsc_HUG env
         let dflags = hsc_dflags env
         opt <- getIdeOptions
-        let getTargetFor modName nfp
-                | Just (TargetFile nfp') <- HM.lookup (TargetFile nfp) targetsMap = do
-                    -- reuse the existing NormalizedFilePath in order to maximize sharing
-                    itExists <- getFileExists nfp'
-                    return $ if itExists then Just nfp' else Nothing
-                | Just tt <- HM.lookup (TargetModule modName) targets = do
-                    -- reuse the existing NormalizedFilePath in order to maximize sharing
-                    let ttmap = HM.mapWithKey const (HashSet.toMap tt)
-                        nfp' = HM.lookupDefault nfp nfp ttmap
-                    itExists <- getFileExists nfp'
-                    return $ if itExists then Just nfp' else Nothing
-                | otherwise = do
-                    itExists <- getFileExists nfp
-                    return $ if itExists then Just nfp else Nothing
+
+        moduleMaps <- use_ GetModulesPaths file
+
         (diags, imports') <- fmap unzip $ forM imports $ \(isSource, (mbPkgName, modName)) -> do
-            diagOrImp <- locateModule (hscSetFlags dflags env) import_dirs (optExtensions opt) getTargetFor modName mbPkgName isSource
+
+            diagOrImp <- locateModule moduleMaps (hscSetFlags dflags env) import_dirs (optExtensions opt) modName mbPkgName isSource
             case diagOrImp of
                 Left diags              -> pure (diags, Just (modName, Nothing))
                 Right (FileImport path) -> pure ([], Just (modName, Just path))
@@ -626,6 +624,108 @@ getModuleGraphRule recorder = defineEarlyCutOffNoFile (cmapWithPrio LogShake rec
   fs <- toKnownFiles <$> useNoFile_ GetKnownTargets
   dependencyInfoForFiles (HashSet.toList fs)
 
+getModulesPathsRule :: Recorder (WithPriority Log) -> Rules ()
+getModulesPathsRule recorder = defineEarlyCutoff (cmapWithPrio LogShake recorder) $ Rule $ \GetModulesPaths file -> do
+  env_eq <- use_ GhcSession file
+
+  ShakeExtras{moduleToPathCache} <- getShakeExtras
+
+  cache <- liftIO (readTVarIO moduleToPathCache)
+  case Map.lookup (envUnique env_eq) cache of
+    Just res -> pure (mempty, ([], Just res))
+    Nothing -> do
+      let env = hscEnv env_eq
+      let import_dirs = map (second homeUnitEnv_dflags) $ hugElts $ hsc_HUG env
+      opt <- getIdeOptions
+      let exts = (optExtensions opt)
+      let acceptedExtensions = concatMap (\x -> ['.':x, '.':x <> "-boot"]) exts
+
+      (unzip -> (a, b)) <- flip mapM import_dirs $ \(u, dyn) -> do
+        (unzip -> (a, b)) <- flip mapM (importPaths dyn) $ \dir' -> do
+          let dir = normalise dir'
+          let predicate path = pure (equalFilePath path dir || case takeFileName path of
+               [] -> False
+               (x:_) -> isUpper x)
+          let dir_number_directories = length (splitDirectories dir)
+          let toModule file = mkModuleName (intercalate "." $ drop dir_number_directories (splitDirectories (dropExtension file)))
+
+          -- TODO: we are taking/dropping extension, this could be factorized to save a few cpu cycles ;)
+          -- TODO: do acceptedextensions needs to be a set ? or a vector?
+          -- If the directory is empty, we return an empty list of modules
+          -- using 'catch' instead of an exception which would kill the LSP
+          modules <- (fmap (\path -> (toModule path, toNormalizedFilePath' path)) . filter (\y -> takeExtension y `elem` acceptedExtensions) <$> liftIO (listFilesInside predicate dir))
+                 `catch` (\(_ :: IOException) -> pure [])
+          let isSourceModule (_, path) = "-boot" `isSuffixOf` fromNormalizedFilePath path
+          let (sourceModules, notSourceModules) = partition isSourceModule modules
+          pure $ (Map.fromList notSourceModules, Map.fromList sourceModules)
+        pure (fmap (u,) $ mconcat a, fmap (u, ) $ mconcat b)
+
+      let res = (mconcat a, mconcat b)
+      -- Extend the current module map with all the known targets
+      resExtended <- extendModuleMapWithKnownTargets file res
+
+      liftIO $ atomically $ modifyTVar' moduleToPathCache (Map.insert (envUnique env_eq) resExtended)
+
+      pure (mempty, ([], Just resExtended))
+
+
+-- | Extend the map from module name to filepath (exiting on the drive) with
+-- the list of known targets provided by HLS
+--
+-- These known targets are files which were recently created and not yet saved
+-- to the filesystem.
+--
+-- This code is not really efficient (O(import_dirs * known_targets)), which
+-- can quickly represents millions of tests. It could be improved, but
+-- everything only happen once everytime known_targets is updated (which only
+-- happen when the project start or when a file is added) and the code is not
+-- doing any IO.
+extendModuleMapWithKnownTargets
+    :: NormalizedFilePath -> (Map.Map ModuleName (UnitId, NormalizedFilePath), Map.Map ModuleName (UnitId, NormalizedFilePath)) ->
+       Action (Map.Map ModuleName (UnitId, NormalizedFilePath), Map.Map ModuleName (UnitId, NormalizedFilePath))
+extendModuleMapWithKnownTargets file (notSourceModules, sourceModules) = do
+  KnownTargets targetsMap <- useNoFile_ GetKnownTargets
+
+  env_eq <- use_ GhcSession file
+  let env = hscEnv env_eq
+  let import_dirs = map (second homeUnitEnv_dflags) $ hugElts $ hsc_HUG env
+  opt <- getIdeOptions
+  let exts = (optExtensions opt)
+  let acceptedExtensions = concatMap (\x -> ['.':x, '.':x <> "-boot"]) exts
+
+  let (unzip -> (catMaybes -> a, catMaybes -> b)) = do
+        (u, dyn) <- import_dirs
+        -- TODO: avoid using so much `FilePath` logic AND please please,
+        -- normalize earlier.
+        --
+        -- The normalise here is in order to remove the trailing `.` which
+        -- could break the comparison later.
+        (normalise -> dir') <- importPaths dyn
+        let dirComponents = splitDirectories dir'
+        let dir_number_directories = length dirComponents
+        -- TODO: the _target may represents something different than the path
+        -- stored in paths. This need to be investigated.
+        (_target, paths) <- HM.toList targetsMap
+        path <- HashSet.toList paths
+        let pathString = fromNormalizedFilePath path
+        let pathComponents = splitDirectories pathString
+
+        -- Ensure this file is in the directory
+        guard $ dirComponents `isPrefixOf` pathComponents
+
+        -- Ensure that this extension is accepted
+        guard $ takeExtension pathString `elem` acceptedExtensions
+        let modName = mkModuleName (intercalate "." $ drop dir_number_directories (splitDirectories (dropExtension pathString)))
+        let isSourceModule = "-boot" `isSuffixOf` pathString
+        if isSourceModule
+        then
+           pure (Nothing, Just (modName, (u, path)))
+        else
+           pure (Just (modName, (u, path)), Nothing)
+
+  pure $ (Map.fromList a <> notSourceModules, Map.fromList b <> sourceModules)
+
+
 dependencyInfoForFiles :: [NormalizedFilePath] -> Action (BS.ByteString, DependencyInformation)
 dependencyInfoForFiles fs = do
   (rawDepInfo, bm) <- rawDependencyInformation fs
@@ -708,6 +808,7 @@ loadGhcSession recorder ghcSessionDepsConfig = do
         IdeGhcSession{loadSessionFun} <- useNoFile_ GhcSessionIO
         -- loading is always returning a absolute path now
         (val,deps) <- liftIO $ loadSessionFun $ fromNormalizedFilePath file
+        -- TODO: this is responsible for a LOT of allocations
 
         -- add the deps to the Shake graph
         let addDependency fp = do
@@ -1231,6 +1332,7 @@ mainRule recorder RulesConfig{..} = do
     getModIfaceRule recorder
     getModSummaryRule templateHaskellWarning recorder
     getModuleGraphRule recorder
+    getModulesPathsRule recorder
     getFileHashRule recorder
     knownFilesRule recorder
     getClientSettingsRule recorder

ghcide/src/Development/IDE/Core/Shake.hs

@@ -131,11 +131,12 @@ import           Development.IDE.Types.Options          as Options
 import qualified Language.LSP.Protocol.Message          as LSP
 import qualified Language.LSP.Server                    as LSP
 
+import           Data.Map.Strict                        (Map)
 import           Development.IDE.Core.Tracing
 import           Development.IDE.Core.WorkerThread
-import           Development.IDE.GHC.Compat             (NameCache,
+import           Development.IDE.GHC.Compat             (ModuleName, NameCache,
                                                          NameCacheUpdater,
-                                                         initNameCache,
+                                                         UnitId, initNameCache,
                                                          knownKeyNames)
 import           Development.IDE.GHC.Orphans            ()
 import           Development.IDE.Graph                  hiding (ShakeValue,
@@ -311,6 +312,10 @@ data ShakeExtras = ShakeExtras
     ,ideNc :: NameCache
     -- | A mapping of module name to known target (or candidate targets, if missing)
     ,knownTargetsVar :: TVar (Hashed KnownTargets)
+    ,moduleToPathCache :: TVar (Map
+     Unique
+     (Map ModuleName (UnitId, NormalizedFilePath),
+      Map ModuleName (UnitId, NormalizedFilePath)))
     -- | A mapping of exported identifiers for local modules. Updated on kick
     ,exportsMap :: TVar ExportsMap
     -- | A work queue for actions added via 'runInShakeSession'
@@ -705,6 +710,8 @@ shakeOpen recorder lspEnv defaultConfig idePlugins debouncer
         dirtyKeys <- newTVarIO mempty
         -- Take one VFS snapshot at the start
         vfsVar <- newTVarIO =<< vfsSnapshot lspEnv
+
+        moduleToPathCache <- newTVarIO mempty
         pure ShakeExtras{shakeRecorder = recorder, ..}
     shakeDb  <-
         shakeNewDatabase
@@ -1482,3 +1489,5 @@ runWithSignal msgStart msgEnd files rule = do
   kickSignal testing lspEnv files msgStart
   void $ uses rule files
   kickSignal testing lspEnv files msgEnd
+
+