Use Data.HashMap.lookupKey and Data.HashSet.lookupElement (#4802)

Author: sjakobi

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

@@ -321,7 +321,7 @@ 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
+        (KnownTargets targets) <- useNoFile_ GetKnownTargets
 #if MIN_VERSION_ghc(9,13,0)
         let imports = [(False, lvl, mbPkgName, modName) | (lvl, mbPkgName, modName) <- ms_textual_imps ms]
                    ++ [(True, NormalLevel, NoPkgQual, noLoc modName) | L _ modName <- ms_srcimps ms]
@@ -334,14 +334,13 @@ getLocatedImportsRule recorder =
         let dflags = hsc_dflags env
         opt <- getIdeOptions
         let getTargetFor modName nfp
-                | Just (TargetFile nfp') <- HM.lookup (TargetFile nfp) targetsMap = do
+                | Just (TargetFile nfp') <- HM.lookupKey (TargetFile nfp) targets = 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
+                    let nfp' = fromMaybe nfp $ HashSet.lookupElement nfp tt
                     itExists <- getFileExists nfp'
                     return $ if itExists then Just nfp' else Nothing
                 | otherwise = do

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

@@ -19,49 +19,29 @@ import           Development.IDE.Types.Location
 import           GHC.Generics
 
 -- | A mapping of module name to known files
-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
+newtype KnownTargets = KnownTargets
+  { targetMap :: (HashMap Target (HashSet NormalizedFilePath)) }
+  deriving Show
 
 
 unionKnownTargets :: KnownTargets -> KnownTargets -> KnownTargets
-unionKnownTargets (KnownTargets tm nm) (KnownTargets tm' nm') =
-  KnownTargets (HMap.unionWith (<>) tm tm') (HMap.union nm nm')
+unionKnownTargets (KnownTargets tm) (KnownTargets tm') =
+  KnownTargets (HMap.unionWith (<>) tm tm')
 
 mkKnownTargets :: [(Target, HashSet NormalizedFilePath)] -> KnownTargets
-mkKnownTargets vs = KnownTargets (HMap.fromList vs) (HMap.fromList [(k,k) | (k,_) <- vs ])
+mkKnownTargets vs = KnownTargets (HMap.fromList vs)
 
 instance NFData KnownTargets where
-  rnf (KnownTargets tm nm) = rnf tm `seq` rnf nm `seq` ()
+  rnf (KnownTargets tm) = rnf tm `seq` ()
 
 instance Eq KnownTargets where
   k1 == k2 = targetMap k1 == targetMap k2
 
 instance Hashable KnownTargets where
-  hashWithSalt s (KnownTargets hm _) = hashWithSalt s hm
+  hashWithSalt s (KnownTargets hm) = hashWithSalt s hm
 
 emptyKnownTargets :: KnownTargets
-emptyKnownTargets = KnownTargets HMap.empty HMap.empty
+emptyKnownTargets = KnownTargets HMap.empty
 
 data Target = TargetModule ModuleName | TargetFile NormalizedFilePath
   deriving ( Eq, Ord, Generic, Show )