Make module membership on ModuleGraph faster

When loading/reloading with a large number of modules
(>5000) the cost of linear lookups becomes significant.

The changes here made `:reload` go from 6s to 1s on my
test case.

The bottlenecks were `needsLinker` in `DriverPipeline` and
`getModLoop` in `GhcMake`.

Test Plan: ./validate

Reviewers: simonmar, austin, bgamari

Subscribers: thomie, rwbarton

Differential Revision: https://phabricator.haskell.org/D3703

8 files changed, 132 insertions, 53 deletions

diff --git a/compiler/backpack/DriverBkp.hs b/compiler/backpack/DriverBkp.hs
index 6123bc8133..4324e5763b 100644
--- a/compiler/backpack/DriverBkp.hs
+++ b/compiler/backpack/DriverBkp.hs
@@ -288,7 +288,8 @@ buildUnit session cid insts lunit = do
         let hi_dir = expectJust (panic "hiDir Backpack") $ hiDir dflags
             export_mod ms = (ms_mod_name ms, ms_mod ms)
             -- Export everything!
-            mods = [ export_mod ms | ms <- mod_graph, ms_hsc_src ms == HsSrcFile ]
+            mods = [ export_mod ms | ms <- mgModSummaries mod_graph
+                                   , ms_hsc_src ms == HsSrcFile ]
 
         -- Compile relevant only
         hsc_env <- getSession
@@ -660,7 +661,7 @@ hsunitModuleGraph dflags unit = do
             else fmap Just $ summariseRequirement pn mod_name
 
     -- 3. Return the kaboodle
-    return (nodes ++ req_nodes)
+    return $ mkModuleGraph $ nodes ++ req_nodes
 
 summariseRequirement :: PackageName -> ModuleName -> BkpM ModSummary
 summariseRequirement pn mod_name = do
diff --git a/compiler/basicTypes/Module.hs b/compiler/basicTypes/Module.hs
index c693e7a0f4..ab1f391e04 100644
--- a/compiler/basicTypes/Module.hs
+++ b/compiler/basicTypes/Module.hs
@@ -132,7 +132,7 @@ module Module
         -- * Sets of Modules
         ModuleSet,
         emptyModuleSet, mkModuleSet, moduleSetElts,
-        extendModuleSet, extendModuleSetList,
+        extendModuleSet, extendModuleSetList, delModuleSet,
         elemModuleSet, intersectModuleSet, minusModuleSet, unionModuleSet,
         unitModuleSet
     ) where
@@ -1276,6 +1276,9 @@ intersectModuleSet = coerce Set.intersection
 minusModuleSet :: ModuleSet -> ModuleSet -> ModuleSet
 minusModuleSet = coerce Set.difference
 
+delModuleSet :: ModuleSet -> Module -> ModuleSet
+delModuleSet = coerce (flip Set.delete)
+
 unionModuleSet :: ModuleSet -> ModuleSet -> ModuleSet
 unionModuleSet = coerce Set.union
 
diff --git a/compiler/main/DriverMkDepend.hs b/compiler/main/DriverMkDepend.hs
index dc18a31174..8cf14c57e5 100644
--- a/compiler/main/DriverMkDepend.hs
+++ b/compiler/main/DriverMkDepend.hs
@@ -75,11 +75,11 @@ doMkDependHS srcs = do
     targets <- mapM (\s -> GHC.guessTarget s Nothing) srcs
     GHC.setTargets targets
     let excl_mods = depExcludeMods dflags
-    mod_summaries <- GHC.depanal excl_mods True {- Allow dup roots -}
+    module_graph <- GHC.depanal excl_mods True {- Allow dup roots -}
 
     -- Sort into dependency order
     -- There should be no cycles
-    let sorted = GHC.topSortModuleGraph False mod_summaries Nothing
+    let sorted = GHC.topSortModuleGraph False module_graph Nothing
 
     -- Print out the dependencies if wanted
     liftIO $ debugTraceMsg dflags 2 (text "Module dependencies" $$ ppr sorted)
@@ -91,7 +91,7 @@ doMkDependHS srcs = do
     mapM_ (liftIO . processDeps dflags hsc_env excl_mods root (mkd_tmp_hdl files)) sorted
 
     -- If -ddump-mod-cycles, show cycles in the module graph
-    liftIO $ dumpModCycles dflags mod_summaries
+    liftIO $ dumpModCycles dflags module_graph
 
     -- Tidy up
     liftIO $ endMkDependHS dflags files
@@ -338,8 +338,8 @@ endMkDependHS dflags
 --              Module cycles
 -----------------------------------------------------------------
 
-dumpModCycles :: DynFlags -> [ModSummary] -> IO ()
-dumpModCycles dflags mod_summaries
+dumpModCycles :: DynFlags -> ModuleGraph -> IO ()
+dumpModCycles dflags module_graph
   | not (dopt Opt_D_dump_mod_cycles dflags)
   = return ()
 
@@ -351,7 +351,8 @@ dumpModCycles dflags mod_summaries
   where
 
     cycles :: [[ModSummary]]
-    cycles = [ c | CyclicSCC c <- GHC.topSortModuleGraph True mod_summaries Nothing ]
+    cycles =
+      [ c | CyclicSCC c <- GHC.topSortModuleGraph True module_graph Nothing ]
 
     pp_cycles = vcat [ (text "---------- Cycle" <+> int n <+> ptext (sLit "----------"))
                         $$ pprCycle c $$ blankLine
@@ -379,7 +380,8 @@ pprCycle summaries = pp_group (CyclicSCC summaries)
 
           loop_breaker = head boot_only
           all_others   = tail boot_only ++ others
-          groups = GHC.topSortModuleGraph True all_others Nothing
+          groups =
+            GHC.topSortModuleGraph True (mkModuleGraph all_others) Nothing
 
     pp_ms summary = text mod_str <> text (take (20 - length mod_str) (repeat ' '))
                        <+> (pp_imps empty (map snd (ms_imps summary)) $$
diff --git a/compiler/main/DriverPipeline.hs b/compiler/main/DriverPipeline.hs
index a6873fb640..3fc35e5992 100644
--- a/compiler/main/DriverPipeline.hs
+++ b/compiler/main/DriverPipeline.hs
@@ -236,10 +236,7 @@ compileOne' m_tc_result mHscMessage
        input_fn    = expectJust "compile:hs" (ml_hs_file location)
        input_fnpp  = ms_hspp_file summary
        mod_graph   = hsc_mod_graph hsc_env0
-       needsLinker = any (\ModSummary {ms_hspp_opts} ->
-                            xopt LangExt.TemplateHaskell ms_hspp_opts
-                            || xopt LangExt.QuasiQuotes ms_hspp_opts
-                         ) mod_graph
+       needsLinker = needsTemplateHaskellOrQQ mod_graph
        isDynWay    = any (== WayDyn) (ways dflags0)
        isProfWay   = any (== WayProf) (ways dflags0)
        internalInterpreter = not (gopt Opt_ExternalInterpreter dflags0)
diff --git a/compiler/main/GHC.hs b/compiler/main/GHC.hs
index 4a45bea2e0..3ca07f1443 100644
--- a/compiler/main/GHC.hs
+++ b/compiler/main/GHC.hs
@@ -59,7 +59,8 @@ module GHC (
         compileToCoreModule, compileToCoreSimplified,
 
         -- * Inspecting the module structure of the program
-        ModuleGraph, emptyMG, mapMG,
+        ModuleGraph, emptyMG, mapMG, mkModuleGraph, mgModSummaries,
+        mgLookupModule,
         ModSummary(..), ms_mod_name, ModLocation(..),
         getModSummary,
         getModuleGraph,
@@ -873,7 +874,10 @@ type TypecheckedSource = LHsBinds GhcTc
 getModSummary :: GhcMonad m => ModuleName -> m ModSummary
 getModSummary mod = do
    mg <- liftM hsc_mod_graph getSession
-   case [ ms | ms <- mg, ms_mod_name ms == mod, not (isBootSummary ms) ] of
+   let mods_by_name = [ ms | ms <- mgModSummaries mg
+                      , ms_mod_name ms == mod
+                      , not (isBootSummary ms) ]
+   case mods_by_name of
      [] -> do dflags <- getDynFlags
               liftIO $ throwIO $ mkApiErr dflags (text "Module not part of module graph")
      [ms] -> return ms
@@ -1023,7 +1027,7 @@ compileCore simplify fn = do
    _ <- load LoadAllTargets
    -- Then find dependencies
    modGraph <- depanal [] True
-   case find ((== fn) . msHsFilePath) modGraph of
+   case find ((== fn) . msHsFilePath) (mgModSummaries modGraph) of
      Just modSummary -> do
        -- Now we have the module name;
        -- parse, typecheck and desugar the module
@@ -1111,7 +1115,7 @@ data ModuleInfo = ModuleInfo {
 getModuleInfo :: GhcMonad m => Module -> m (Maybe ModuleInfo)  -- XXX: Maybe X
 getModuleInfo mdl = withSession $ \hsc_env -> do
   let mg = hsc_mod_graph hsc_env
-  if mdl `elem` map ms_mod mg
+  if mgElemModule mg mdl
         then liftIO $ getHomeModuleInfo hsc_env mdl
         else do
   {- if isHomeModule (hsc_dflags hsc_env) mdl
diff --git a/compiler/main/GhcMake.hs b/compiler/main/GhcMake.hs
index f4ea4de28d..f4a9a319ac 100644
--- a/compiler/main/GhcMake.hs
+++ b/compiler/main/GhcMake.hs
@@ -138,9 +138,11 @@ depanal excluded_mods allow_dup_roots = do
     -- cached finder data.
     liftIO $ flushFinderCaches hsc_env
 
-    mod_graphE <- liftIO $ downsweep hsc_env old_graph
+    mod_summariesE <- liftIO $ downsweep hsc_env (mgModSummaries old_graph)
                                      excluded_mods allow_dup_roots
-    mod_graph <- reportImportErrors mod_graphE
+    mod_summaries <- reportImportErrors mod_summariesE
+
+    let mod_graph = mkModuleGraph mod_summaries
 
     warnMissingHomeModules hsc_env mod_graph
 
@@ -193,7 +195,7 @@ warnMissingHomeModules hsc_env mod_graph =
     is_my_target _ _ = False
 
     missing = map (moduleName . ms_mod) $
-      filter (not . is_known_module) mod_graph
+      filter (not . is_known_module) (mgModSummaries mod_graph)
 
     msg
       | gopt Opt_BuildingCabalPackage dflags
@@ -253,7 +255,7 @@ load' how_much mHscMessage mod_graph = do
     -- (see msDeps)
     let all_home_mods =
           mkUniqSet [ ms_mod_name s
-                    | s <- mod_graph, not (isBootSummary s)]
+                    | s <- mgModSummaries mod_graph, not (isBootSummary s)]
     -- TODO: Figure out what the correct form of this assert is. It's violated
     -- when you have HsBootMerge nodes in the graph: then you'll have hs-boot
     -- files without corresponding hs files.
@@ -422,7 +424,7 @@ load' how_much mHscMessage mod_graph = do
           let no_hs_main = gopt Opt_NoHsMain dflags
           let
             main_mod = mainModIs dflags
-            a_root_is_Main = any ((==main_mod).ms_mod) mod_graph
+            a_root_is_Main = mgElemModule mod_graph main_mod
             do_linking = a_root_is_Main || no_hs_main || ghcLink dflags == LinkDynLib || ghcLink dflags == LinkStaticLib
 
           -- link everything together
@@ -543,8 +545,7 @@ guessOutputFile = modifySession $ \env ->
         !mod_graph = hsc_mod_graph env
         mainModuleSrcPath :: Maybe String
         mainModuleSrcPath = do
-            let isMain = (== mainModIs dflags) . ms_mod
-            [ms] <- return (filter isMain mod_graph)
+            ms <- mgLookupModule mod_graph (mainModIs dflags)
             ml_hs_file (ms_location ms)
         name = fmap dropExtension mainModuleSrcPath
 
@@ -889,13 +890,19 @@ parUpsweep n_jobs mHscMessage old_hpt stable_mods cleanup sccs = do
     -- The list of all loops in the compilation graph.
     -- NB: For convenience, the last module of each loop (aka the module that
     -- finishes the loop) is prepended to the beginning of the loop.
-    let comp_graph_loops = go (map fstOf3 (reverse comp_graph))
+    let graph = map fstOf3 (reverse comp_graph)
+        boot_modules = mkModuleSet [ms_mod ms | ms <- graph, isBootSummary ms]
+        comp_graph_loops = go graph boot_modules
           where
-            go [] = []
-            go (ms:mss) | Just loop <- getModLoop ms (ms:mss)
-                        = map mkBuildModule (ms:loop) : go mss
-                        | otherwise
-                        = go mss
+            remove ms bm
+              | isBootSummary ms = delModuleSet bm (ms_mod ms)
+              | otherwise = bm
+            go [] _ = []
+            go mg@(ms:mss) boot_modules
+              | Just loop <- getModLoop ms mg (`elemModuleSet` boot_modules)
+              = map mkBuildModule (ms:loop) : go mss (remove ms boot_modules)
+              | otherwise
+              = go mss (remove ms boot_modules)
 
     -- Build a Map out of the compilation graph with which we can efficiently
     -- look up the result MVar associated with a particular home module.
@@ -1236,12 +1243,22 @@ upsweep
 
 upsweep mHscMessage old_hpt stable_mods cleanup sccs = do
    dflags <- getSessionDynFlags
-   (res, done) <- upsweep' old_hpt [] sccs 1 (length sccs)
+   (res, done) <- upsweep' old_hpt emptyMG sccs 1 (length sccs)
                            (unitIdsToCheck dflags) done_holes
-   return (res, reverse done)
+   return (res, reverse $ mgModSummaries done)
  where
   done_holes = emptyUniqSet
 
+  upsweep'
+    :: GhcMonad m
+    => HomePackageTable
+    -> ModuleGraph
+    -> [SCC ModSummary]
+    -> Int
+    -> Int
+    -> [UnitId]
+    -> UniqSet ModuleName
+    -> m (SuccessFlag, ModuleGraph)
   upsweep' _old_hpt done
      [] _ _ uids_to_check _
    = do hsc_env <- getSession
@@ -1319,7 +1336,7 @@ upsweep mHscMessage old_hpt stable_mods cleanup sccs = do
                     old_hpt1 | isBootSummary mod = old_hpt
                              | otherwise = delFromHpt old_hpt this_mod
 
-                    done' = mod:done
+                    done' = extendMG done mod
 
                         -- fixup our HomePackageTable after we've finished compiling
                         -- a mutually-recursive loop.  We have to do this again
@@ -1643,7 +1660,7 @@ Following this fix, GHC can compile itself with --make -O2.
 
 reTypecheckLoop :: HscEnv -> ModSummary -> ModuleGraph -> IO HscEnv
 reTypecheckLoop hsc_env ms graph
-  | Just loop <- getModLoop ms graph
+  | Just loop <- getModLoop ms mss appearsAsBoot
   -- SOME hs-boot files should still
   -- get used, just not the loop-closer.
   , let non_boot = filter (\l -> not (isBootSummary l &&
@@ -1651,11 +1668,18 @@ reTypecheckLoop hsc_env ms graph
   = typecheckLoop (hsc_dflags hsc_env) hsc_env (map ms_mod_name non_boot)
   | otherwise
   = return hsc_env
-
-getModLoop :: ModSummary -> ModuleGraph -> Maybe [ModSummary]
-getModLoop ms graph
+  where
+  mss = mgModSummaries graph
+  appearsAsBoot = (`elemModuleSet` mgBootModules graph)
+
+getModLoop
+  :: ModSummary
+  -> [ModSummary]
+  -> (Module -> Bool) -- check if a module appears as a boot module in 'graph'
+  -> Maybe [ModSummary]
+getModLoop ms graph appearsAsBoot
   | not (isBootSummary ms)
-  , any (\m -> ms_mod m == this_mod && isBootSummary m) graph
+  , appearsAsBoot this_mod
   , let mss = reachableBackwards (ms_mod_name ms) graph
   = Just mss
   | otherwise
@@ -1694,7 +1718,7 @@ reachableBackwards mod summaries
 topSortModuleGraph
           :: Bool
           -- ^ Drop hi-boot nodes? (see below)
-          -> [ModSummary]
+          -> ModuleGraph
           -> Maybe ModuleName
              -- ^ Root module name.  If @Nothing@, use the full graph.
           -> [SCC ModSummary]
@@ -1713,9 +1737,10 @@ topSortModuleGraph
 --              the a source-import of Foo is an import of Foo
 --              The resulting graph has no hi-boot nodes, but can be cyclic
 
-topSortModuleGraph drop_hs_boot_nodes summaries mb_root_mod
+topSortModuleGraph drop_hs_boot_nodes module_graph mb_root_mod
   = map (fmap summaryNodeSummary) $ stronglyConnCompG initial_graph
   where
+    summaries = mgModSummaries module_graph
     -- stronglyConnCompG flips the original order, so if we reverse
     -- the summaries we get a stable topological sort.
     (graph, lookup_node) =
@@ -1999,7 +2024,7 @@ enableCodeGenForTH target nodemap =
       [ ms
       | mss <- Map.elems nodemap
       , Right ms <- mss
-      , needsTemplateHaskellOrQQ $ [ms]
+      , isTemplateHaskellOrQQNonBoot ms
       ]
 
     -- find the set of all transitive dependencies of a list of modules.
diff --git a/compiler/main/HscMain.hs b/compiler/main/HscMain.hs
index 196e309caa..c514e5b017 100644
--- a/compiler/main/HscMain.hs
+++ b/compiler/main/HscMain.hs
@@ -184,7 +184,7 @@ newHscEnv dflags = do
     iserv_mvar <- newMVar Nothing
     return HscEnv {  hsc_dflags       = dflags
                   ,  hsc_targets      = []
-                  ,  hsc_mod_graph    = []
+                  ,  hsc_mod_graph    = emptyMG
                   ,  hsc_IC           = emptyInteractiveContext dflags
                   ,  hsc_HPT          = emptyHomePackageTable
                   ,  hsc_EPS          = eps_var
diff --git a/compiler/main/HscTypes.hs b/compiler/main/HscTypes.hs
index f7a8140583..e064147965 100644
--- a/compiler/main/HscTypes.hs
+++ b/compiler/main/HscTypes.hs
@@ -5,6 +5,7 @@
 -}
 
 {-# LANGUAGE CPP, ScopedTypeVariables #-}
+{-# LANGUAGE RecordWildCards #-}
 
 -- | Types for the per-module compiler
 module HscTypes (
@@ -12,11 +13,14 @@ module HscTypes (
         HscEnv(..), hscEPS,
         FinderCache, FindResult(..), InstalledFindResult(..),
         Target(..), TargetId(..), pprTarget, pprTargetId,
-        needsTemplateHaskellOrQQ,
-        ModuleGraph, emptyMG, mapMG,
         HscStatus(..),
         IServ(..),
 
+        -- * ModuleGraph
+        ModuleGraph, emptyMG, mkModuleGraph, extendMG, mapMG,
+        mgModSummaries, mgElemModule, mgLookupModule,
+        needsTemplateHaskellOrQQ, mgBootModules,
+
         -- * Hsc monad
         Hsc(..), runHsc, runInteractiveHsc,
 
@@ -28,7 +32,7 @@ module HscTypes (
 
         ModSummary(..), ms_imps, ms_installed_mod, ms_mod_name, showModMsg, isBootSummary,
         msHsFilePath, msHiFilePath, msObjFilePath,
-        SourceModified(..),
+        SourceModified(..), isTemplateHaskellOrQQNonBoot,
 
         -- * Information about the module being compiled
         -- (re-exported from DriverPhases)
@@ -2618,8 +2622,16 @@ soExt platform
 --
 -- The graph is not necessarily stored in topologically-sorted order.  Use
 -- 'GHC.topSortModuleGraph' and 'Digraph.flattenSCC' to achieve this.
-type ModuleGraph = [ModSummary]
-
+data ModuleGraph = ModuleGraph
+  { mg_mss :: [ModSummary]
+  , mg_non_boot :: ModuleEnv ModSummary
+    -- a map of all non-boot ModSummaries keyed by Modules
+  , mg_boot :: ModuleSet
+    -- a set of boot Modules
+  , mg_needs_th_or_qq :: !Bool
+    -- does any of the modules in mg_mss require TemplateHaskell or
+    -- QuasiQuotes?
+  }
 
 -- | Determines whether a set of modules requires Template Haskell or
 -- Quasi Quotes
@@ -2628,13 +2640,31 @@ type ModuleGraph = [ModSummary]
 -- 'depanal' was called, then each module in the returned module graph will
 -- have Template Haskell enabled whether it is actually needed or not.
 needsTemplateHaskellOrQQ :: ModuleGraph -> Bool
-needsTemplateHaskellOrQQ mg = any isTemplateHaskellOrQQNonBoot mg
-
-emptyMG :: ModuleGraph
-emptyMG = []
+needsTemplateHaskellOrQQ mg = mg_needs_th_or_qq mg
 
+-- | Map a function 'f' over all the 'ModSummaries'.
+-- To preserve invariants 'f' can't change the isBoot status.
 mapMG :: (ModSummary -> ModSummary) -> ModuleGraph -> ModuleGraph
-mapMG = map
+mapMG f mg@ModuleGraph{..} = mg
+  { mg_mss = map f mg_mss
+  , mg_non_boot = mapModuleEnv f mg_non_boot
+  }
+
+mgBootModules :: ModuleGraph -> ModuleSet
+mgBootModules ModuleGraph{..} = mg_boot
+
+mgModSummaries :: ModuleGraph -> [ModSummary]
+mgModSummaries = mg_mss
+
+mgElemModule :: ModuleGraph -> Module -> Bool
+mgElemModule ModuleGraph{..} m = elemModuleEnv m mg_non_boot
+
+-- | Look up a ModSummary in the ModuleGraph
+mgLookupModule :: ModuleGraph -> Module -> Maybe ModSummary
+mgLookupModule ModuleGraph{..} m = lookupModuleEnv mg_non_boot m
+
+emptyMG :: ModuleGraph
+emptyMG = ModuleGraph [] emptyModuleEnv emptyModuleSet False
 
 isTemplateHaskellOrQQNonBoot :: ModSummary -> Bool
 isTemplateHaskellOrQQNonBoot ms =
@@ -2642,6 +2672,23 @@ isTemplateHaskellOrQQNonBoot ms =
     || xopt LangExt.QuasiQuotes (ms_hspp_opts ms)) &&
   not (isBootSummary ms)
 
+-- | Add a ModSummary to ModuleGraph. Assumes that the new ModSummary is
+-- not an element of the ModuleGraph.
+extendMG :: ModuleGraph -> ModSummary -> ModuleGraph
+extendMG ModuleGraph{..} ms = ModuleGraph
+  { mg_mss = ms:mg_mss
+  , mg_non_boot = if isBootSummary ms
+      then mg_non_boot
+      else extendModuleEnv mg_non_boot (ms_mod ms) ms
+  , mg_boot = if isBootSummary ms
+      then extendModuleSet mg_boot (ms_mod ms)
+      else mg_boot
+  , mg_needs_th_or_qq = mg_needs_th_or_qq || isTemplateHaskellOrQQNonBoot ms
+  }
+
+mkModuleGraph :: [ModSummary] -> ModuleGraph
+mkModuleGraph = foldr (flip extendMG) emptyMG
+
 -- | A single node in a 'ModuleGraph'. The nodes of the module graph
 -- are one of:
 --