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: rwbarton, thomie

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

7 files changed, 174 insertions, 54 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/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 eed66b22c1..850f67d061 100644
--- a/compiler/main/DriverPipeline.hs
+++ b/compiler/main/DriverPipeline.hs
@@ -235,10 +235,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 2102009019..3ca07f1443 100644
--- a/compiler/main/GHC.hs
+++ b/compiler/main/GHC.hs
@@ -23,7 +23,7 @@ module GHC (
         gcatch, gbracket, gfinally,
         printException,
         handleSourceError,
-        needsTemplateHaskell,
+        needsTemplateHaskellOrQQ,
 
         -- * Flags and settings
         DynFlags(..), GeneralFlag(..), Severity(..), HscTarget(..), gopt,
@@ -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
@@ -1075,15 +1079,6 @@ compileCore simplify fn = do
 getModuleGraph :: GhcMonad m => m ModuleGraph -- ToDo: DiGraph ModSummary
 getModuleGraph = liftM hsc_mod_graph getSession
 
--- | Determines whether a set of modules requires Template Haskell.
---
--- Note that if the session's 'DynFlags' enabled Template Haskell when
--- 'depanal' was called, then each module in the returned module graph will
--- have Template Haskell enabled whether it is actually needed or not.
-needsTemplateHaskell :: ModuleGraph -> Bool
-needsTemplateHaskell ms =
-    any (xopt LangExt.TemplateHaskell . ms_hspp_opts) ms
-
 -- | Return @True@ <==> module is loaded.
 isLoaded :: GhcMonad m => ModuleName -> m Bool
 isLoaded m = withSession $ \hsc_env ->
@@ -1120,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 134a0607bc..57af356b38 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 = text "Modules are not listed in command line: "
         <> sep (map ppr missing)
@@ -248,7 +250,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.
@@ -417,7 +419,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
@@ -538,8 +540,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
 
@@ -884,13 +885,15 @@ 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 comp_graph_loops = go graph
           where
-            go [] = []
-            go (ms:mss) | Just loop <- getModLoop ms (ms:mss)
-                        = map mkBuildModule (ms:loop) : go mss
-                        | otherwise
-                        = go mss
+            graph = mkModuleGraph (map fstOf3 (reverse comp_graph))
+            go mg
+              | Just (ms, mg') <- mgHead mg =
+                case getModLoop ms mg of
+                 Just loop -> map mkBuildModule (ms:loop) : go mg'
+                 Nothing -> go mg'
+              | otherwise = []
 
     -- Build a Map out of the compilation graph with which we can efficiently
     -- look up the result MVar associated with a particular home module.
@@ -1231,12 +1234,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
@@ -1314,7 +1327,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
@@ -1650,8 +1663,8 @@ reTypecheckLoop hsc_env ms graph
 getModLoop :: ModSummary -> ModuleGraph -> Maybe [ModSummary]
 getModLoop ms graph
   | not (isBootSummary ms)
-  , any (\m -> ms_mod m == this_mod && isBootSummary m) graph
-  , let mss = reachableBackwards (ms_mod_name ms) graph
+  , mgElemBootModule graph this_mod
+  , let mss = reachableBackwards (ms_mod_name ms) (mgModSummaries graph)
   = Just mss
   | otherwise
   = Nothing
@@ -1689,7 +1702,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]
@@ -1708,9 +1721,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) =
diff --git a/compiler/main/HscMain.hs b/compiler/main/HscMain.hs
index b8bd76bedd..8d8e26ebb5 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 fa9c18a3e1..868453d351 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,7 +13,9 @@ module HscTypes (
         HscEnv(..), hscEPS,
         FinderCache, FindResult(..), InstalledFindResult(..),
         Target(..), TargetId(..), pprTarget, pprTargetId,
-        ModuleGraph, emptyMG, mapMG,
+        ModuleGraph, emptyMG, mkModuleGraph, mgHead, mgReverse, extendMG, mapMG,
+        mgModSummaries, mgElemModule, mgElemBootModule, mgLookupModule,
+        needsTemplateHaskellOrQQ,
         HscStatus(..),
         IServ(..),
 
@@ -199,6 +202,7 @@ import Platform
 import Util
 import UniqDSet
 import GHC.Serialized   ( Serialized )
+import qualified GHC.LanguageExtensions as LangExt
 
 import Foreign
 import Control.Monad    ( guard, liftM, ap )
@@ -2606,13 +2610,120 @@ soExt platform
 --
 -- The graph is not necessarily stored in topologically-sorted order.  Use
 -- 'GHC.topSortModuleGraph' and 'Digraph.flattenSCC' to achieve this.
-type ModuleGraph = [ModSummary]
-
-emptyMG :: ModuleGraph
-emptyMG = []
+data ModuleGraph = ModuleGraph
+  { mg_mss :: [ModSummary]
+  , mg_non_boot :: ModuleEnv ModSummary
+    -- a map of all non-boot ModSummaries keyed by Modules
+  , mg_boot :: ModuleEnv ModSummary
+    -- a map of all boot ModSummaries keyed by Modules
+  , mg_needs_th_or_qq :: ModuleEnv ModSummary
+    -- all non-boot Modules that need TemplateHaskell or QuasiQuotes
+  }
 
+-- | Determines whether a set of modules requires Template Haskell or
+-- Quasi Quotes
+--
+-- Note that if the session's 'DynFlags' enabled Template Haskell when
+-- '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 = not $ isEmptyModuleEnv $ 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 ModuleGraph{..} =
+  ModuleGraph
+    { mg_mss = map f mg_mss
+    , mg_non_boot = mapModuleEnv f mg_non_boot
+    , mg_boot = mapModuleEnv f mg_boot
+    , mg_needs_th_or_qq = mapModuleEnv f mg_needs_th_or_qq
+    }
+
+mgModSummaries :: ModuleGraph -> [ModSummary]
+mgModSummaries = mg_mss
+
+mgElemModule :: ModuleGraph -> Module -> Bool
+mgElemModule mg m = mgElemNonBootModule mg m || mgElemBootModule mg m
+
+mgElemBootModule :: ModuleGraph -> Module -> Bool
+mgElemBootModule ModuleGraph{..} m = elemModuleEnv m mg_boot
+
+mgElemNonBootModule :: ModuleGraph -> Module -> Bool
+mgElemNonBootModule ModuleGraph{..} m = elemModuleEnv m mg_non_boot
+
+mgLookupNonBootModule :: ModuleGraph -> Module -> Maybe ModSummary
+mgLookupNonBootModule ModuleGraph{..} m = lookupModuleEnv mg_non_boot m
+
+mgLookupBootModule :: ModuleGraph -> Module -> Maybe ModSummary
+mgLookupBootModule ModuleGraph{..} m = lookupModuleEnv mg_boot m
+
+-- | Look up a ModSummary in the ModuleGraph, in non-boot ModSummaries first,
+-- falling back to boot ModSummaries.
+mgLookupModule :: ModuleGraph -> Module -> Maybe ModSummary
+mgLookupModule mg m =
+  case mgLookupNonBootModule mg m of
+    Nothing -> mgLookupBootModule mg m
+    a -> a
+
+emptyMG :: ModuleGraph
+emptyMG =
+  ModuleGraph [] emptyModuleEnv emptyModuleEnv emptyModuleEnv
+
+-- | Reverse the order of elements in the ModuleGraph.
+mgReverse :: ModuleGraph -> ModuleGraph
+mgReverse mg@ModuleGraph { mg_mss = mg_mss } = mg { mg_mss = reverse mg_mss }
+
+isTemplateHaskellOrQQNonBoot :: ModSummary -> Bool
+isTemplateHaskellOrQQNonBoot ms =
+  (xopt LangExt.TemplateHaskell (ms_hspp_opts 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 = extendModEnvIf (not . isBootSummary) mg_non_boot ms
+  , mg_boot = extendModEnvIf isBootSummary mg_boot ms
+  , mg_needs_th_or_qq =
+      extendModEnvIf isTemplateHaskellOrQQNonBoot mg_needs_th_or_qq ms
+  }
+  where
+  extendModEnvIf
+    :: (ModSummary -> Bool)
+    -> ModuleEnv ModSummary
+    -> ModSummary
+    -> ModuleEnv ModSummary
+  extendModEnvIf p me ms
+    | p ms = extendModuleEnv me (ms_mod ms) ms
+    | otherwise = me
+
+-- | Take the first element from the ModuleGraph and remove it from the
+-- graph.
+mgHead :: ModuleGraph -> Maybe (ModSummary, ModuleGraph)
+mgHead ModuleGraph { mg_mss = [] } = Nothing
+mgHead ModuleGraph { mg_mss = (ms:mss), ..} = Just (ms, mg')
+  where
+  mg' = ModuleGraph
+    { mg_mss = mss
+    , mg_non_boot = delModEnvIf (not . isBootSummary) mg_non_boot ms
+    , mg_boot = delModEnvIf isBootSummary mg_boot ms
+    , mg_needs_th_or_qq =
+        delModEnvIf isTemplateHaskellOrQQNonBoot mg_needs_th_or_qq ms
+    }
+  delModEnvIf
+    :: (ModSummary -> Bool)
+    -> ModuleEnv ModSummary
+    -> ModSummary
+    -> ModuleEnv ModSummary
+  delModEnvIf p me ms
+    | p ms = delModuleEnv me (ms_mod ms)
+    | otherwise = me
+
+mkModuleGraph :: [ModSummary] -> ModuleGraph
+mkModuleGraph = foldr (flip extendMG) emptyMG
 
 -- | A single node in a 'ModuleGraph'. The nodes of the module graph
 -- are one of: