Fix quadratic memory usage in GetLocatedImports

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

@@ -505,13 +505,12 @@
                 return [(targetTarget, Set.fromList found)]
           hasUpdate <- atomically $ do
             known <- readTVar knownTargetsVar
-            let known' = flip mapHashed known $ \k ->
-                            HM.unionWith (<>) k $ HM.fromList knownTargets
+            let known' = flip mapHashed known $ \k -> unionKnownTargets k (mkKnownTargets knownTargets)
                 hasUpdate = if known /= known' then Just (unhashed known') else Nothing
             writeTVar knownTargetsVar known'
             pure hasUpdate
           for_ hasUpdate $ \x ->
-            logWith recorder Debug $ LogKnownFilesUpdated x
+            logWith recorder Debug $ LogKnownFilesUpdated (targetMap x)
           return $ toNoFileKey GetKnownTargets
 
     -- Create a new HscEnv from a hieYaml root and a set of options
@@ -660,7 +659,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

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

@@ -321,8 +321,7 @@
 getLocatedImportsRule recorder =
     define (cmapWithPrio LogShake recorder) $ \GetLocatedImports file -> do
         ModSummaryResult{msrModSummary = ms} <- use_ GetModSummaryWithoutTimestamps file
-        targets <- useNoFile_ GetKnownTargets
-        let targetsMap = HM.mapWithKey const targets
+        (KnownTargets targets targetsMap) <- useNoFile_ GetKnownTargets
         let imports = [(False, imp) | imp <- ms_textual_imps ms] ++ [(True, imp) | imp <- ms_srcimps ms]
         env_eq <- use_ GhcSession file
         let env = hscEnvWithImportPaths env_eq
@@ -824,7 +823,7 @@
             { source_version = ver
             , old_value = m_old
             , get_file_version = use GetModificationTime_{missingFileDiagnostics = False}
             , get_linkable_hashes = \fs -> map (snd . fromJust . hirCoreFp) <$> uses_ GetModIface fs
             , regenerate = regenerateHiFile session f ms
             }
       r <- loadInterface (hscEnv session) ms linkableType recompInfo
@@ -1106,7 +1105,7 @@
             -- thus bump its modification time, forcing this rule to be rerun every time.
             exists <- liftIO $ doesFileExist obj_file
             mobj_time <- liftIO $
               if exists
               then Just <$> getModTime obj_file
               else pure Nothing
             case mobj_time of

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

@@ -24,7 +24,7 @@
 module Development.IDE.Core.Shake(
     IdeState, shakeSessionInit, shakeExtras, shakeDb, rootDir,
     ShakeExtras(..), getShakeExtras, getShakeExtrasRules,
-    KnownTargets, Target(..), toKnownFiles,
+    KnownTargets(..), Target(..), toKnownFiles, unionKnownTargets, mkKnownTargets,
     IdeRule, IdeResult,
     GetModificationTime(GetModificationTime, GetModificationTime_, missingFileDiagnostics),
     shakeOpen, shakeShut,
@@ -125,7 +125,7 @@
 import           Development.IDE.Core.RuleTypes
 import           Development.IDE.Core.Tracing
 import           Development.IDE.Core.WorkerThread
 import           Development.IDE.GHC.Compat             (NameCache,
                                                          initNameCache,
                                                          knownKeyNames)
 import           Development.IDE.GHC.Orphans            ()
@@ -691,7 +691,7 @@
         publishedDiagnostics <- STM.newIO
         semanticTokensCache <- STM.newIO
         positionMapping <- STM.newIO
-        knownTargetsVar <- newTVarIO $ hashed HMap.empty
+        knownTargetsVar <- newTVarIO $ hashed emptyKnownTargets
         let restartShakeSession = shakeRestart recorder ideState
         persistentKeys <- newTVarIO mempty
         indexPending <- newTVarIO HMap.empty

ghcide/src/Development/IDE/Plugin/Completions.hs

@@ -180,7 +180,7 @@ getCompletionsLSP ide plId
             pm <- useWithStaleFast GetParsedModule npath
             binds <- fromMaybe (mempty, zeroMapping) <$> useWithStaleFast GetBindings npath
             knownTargets <- liftIO $ runAction  "Completion" ide $ useNoFile GetKnownTargets
-            let localModules = maybe [] Map.keys knownTargets
+            let localModules = maybe [] (Map.keys . targetMap) knownTargets
             let lModules = mempty{importableModules = map toModueNameText localModules}
             -- set up the exports map including both package and project-level identifiers
             packageExportsMapIO <- fmap(envPackageExports . fst) <$> useWithStaleFast GhcSession npath

ghcide/src/Development/IDE/Types/KnownTargets.hs

@@ -1,6 +1,11 @@
 {-# LANGUAGE DeriveAnyClass     #-}
 {-# LANGUAGE DerivingStrategies #-}
-module Development.IDE.Types.KnownTargets (KnownTargets, Target(..), toKnownFiles) where
+module Development.IDE.Types.KnownTargets ( KnownTargets(..)
+                                          , emptyKnownTargets
+                                          , mkKnownTargets
+                                          , unionKnownTargets
+                                          , Target(..)
+                                          , toKnownFiles) where
 
 import           Control.DeepSeq
 import           Data.Hashable
@@ -14,11 +19,53 @@ import           Development.IDE.Types.Location
 import           GHC.Generics
 
 -- | A mapping of module name to known files
-type KnownTargets = HashMap Target (HashSet NormalizedFilePath)
+data KnownTargets = KnownTargets
+  { targetMap      :: !(HashMap Target (HashSet NormalizedFilePath))
+  -- | 'normalisingMap' is a cached copy of `HMap.mapKey const targetMap`
+  --
+  -- At startup 'GetLocatedImports' is called on all known files. Say you have 10000
+  -- modules in your project then this leads to 10000 calls to 'GetLocatedImports'
+  -- running concurrently.
+  --
+  -- In `GetLocatedImports` the known targets are consulted and the targetsMap
+  -- is created by mapping the known targets. This map is used for introducing
+  -- sharing amongst filepaths.  This operation copies a local copy of the `target`
+  --  map which is local to the rule.
+  --
+  -- @
+  -- let targetsMap = HMap.mapWithKey const targets
+  -- @
+  --
+  -- So now each rule has a 'HashMap' of size 10000 held locally to it and depending
+  -- on how the threads are scheduled there will be 10000^2 elements in total
+  -- allocated in 'HashMap's. This used a lot of memory.
+  --
+  -- Solution: Return the 'normalisingMap' in the result of the `GetKnownTargets` rule so it is shared across threads.
+  , normalisingMap :: !(HashMap Target Target) } deriving Show
+
+
+unionKnownTargets :: KnownTargets -> KnownTargets -> KnownTargets
+unionKnownTargets (KnownTargets tm nm) (KnownTargets tm' nm') =
+  KnownTargets (HMap.unionWith (<>) tm tm') (HMap.union nm nm')
+
+mkKnownTargets :: [(Target, HashSet NormalizedFilePath)] -> KnownTargets
+mkKnownTargets vs = KnownTargets (HMap.fromList vs) (HMap.fromList [(k,k) | (k,_) <- vs ])
+
+instance NFData KnownTargets where
+  rnf (KnownTargets tm nm) = rnf tm `seq` rnf nm `seq` ()
+
+instance Eq KnownTargets where
+  k1 == k2 = targetMap k1 == targetMap k2
+
+instance Hashable KnownTargets where
+  hashWithSalt s (KnownTargets hm _) = hashWithSalt s hm
+
+emptyKnownTargets :: KnownTargets
+emptyKnownTargets = KnownTargets HMap.empty HMap.empty
 
 data Target = TargetModule ModuleName | TargetFile NormalizedFilePath
   deriving ( Eq, Ord, Generic, Show )
   deriving anyclass (Hashable, NFData)
 
 toKnownFiles :: KnownTargets -> HashSet NormalizedFilePath
-toKnownFiles = HSet.unions . HMap.elems
+toKnownFiles = HSet.unions . HMap.elems . targetMap