driver: Cache the transitive dependency calculation in ModuleGraph

Two reasons for this change:

1. Avoid computing the transitive dependencies when compiling each
   module, this can save a lot of repeated work.
2. More robust to forthcoming changes to support multiple home units.

3 files changed, 135 insertions, 106 deletions

diff --git a/compiler/GHC/Data/Graph/Directed.hs b/compiler/GHC/Data/Graph/Directed.hs
index 62482bfe30..411b22d919 100644
--- a/compiler/GHC/Data/Graph/Directed.hs
+++ b/compiler/GHC/Data/Graph/Directed.hs
@@ -64,6 +64,7 @@ import GHC.Types.Unique.FM
 import qualified Data.IntMap as IM
 import qualified Data.IntSet as IS
 import qualified Data.Map as M
+import qualified Data.Set as S
 
 {-
 ************************************************************************
@@ -374,8 +375,8 @@ reachablesG graph froms = map (gr_vertex_to_node graph) result
                  reachable (gr_int_graph graph) vs
         vs = [ v | Just v <- map (gr_node_to_vertex graph) froms ]
 
-allReachable :: Ord key => Graph node -> (node -> key) -> M.Map key [key]
-allReachable (Graph g from _) conv =  M.fromList [(conv (from v), IS.foldr (\k vs -> conv (from k) : vs) [] vs) | (v, vs) <- IM.toList int_graph]
+allReachable :: Ord key => Graph node -> (node -> key) -> M.Map key (S.Set key)
+allReachable (Graph g from _) conv =  M.fromList [(conv (from v), IS.foldr (\k vs -> conv (from k) `S.insert` vs) S.empty vs) | (v, vs) <- IM.toList int_graph]
   where
     int_graph = reachableGraph g
 
diff --git a/compiler/GHC/Driver/Make.hs b/compiler/GHC/Driver/Make.hs
index efaf68aeaf..2c86b3c22b 100644
--- a/compiler/GHC/Driver/Make.hs
+++ b/compiler/GHC/Driver/Make.hs
@@ -43,7 +43,7 @@ module GHC.Driver.Make (
         implicitRequirementsShallow,
 
         noModError, cyclicModuleErr,
-        moduleGraphNodes, SummaryNode,
+        SummaryNode,
         IsBootInterface(..), mkNodeKey,
 
         ModNodeMap(..), emptyModNodeMap, modNodeMapElems, modNodeMapLookup, modNodeMapInsert
@@ -1322,80 +1322,17 @@ topSortModules drop_hs_boot_nodes summaries mb_root_mod
                      = throwGhcException (ProgramError "module does not exist")
             in graphFromEdgedVerticesUniq (seq root (reachableG graph root))
 
-type SummaryNode = Node Int ModuleGraphNode
-
-summaryNodeKey :: SummaryNode -> Int
-summaryNodeKey = node_key
-
-summaryNodeSummary :: SummaryNode -> ModuleGraphNode
-summaryNodeSummary = node_payload
-
--- | Collect the immediate dependencies of a ModuleGraphNode,
--- optionally avoiding hs-boot dependencies.
--- If the drop_hs_boot_nodes flag is False, and if this is a .hs and there is
--- an equivalent .hs-boot, add a link from the former to the latter.  This
--- has the effect of detecting bogus cases where the .hs-boot depends on the
--- .hs, by introducing a cycle.  Additionally, it ensures that we will always
--- process the .hs-boot before the .hs, and so the HomePackageTable will always
--- have the most up to date information.
-unfilteredEdges :: Bool -> ModuleGraphNode -> [NodeKey]
-unfilteredEdges drop_hs_boot_nodes = \case
-    InstantiationNode iuid ->
-      NodeKey_Module . flip GWIB NotBoot <$> uniqDSetToList (instUnitHoles iuid)
-    ModuleNode (ExtendedModSummary ms bds) ->
-      [ NodeKey_Unit inst_unit | inst_unit <- bds ] ++
-      (NodeKey_Module . flip GWIB hs_boot_key . unLoc <$> ms_home_srcimps ms) ++
-      [ NodeKey_Module $ GWIB (ms_mod_name ms) IsBoot
-      | not $ drop_hs_boot_nodes || ms_hsc_src ms == HsBootFile
-      ] ++
-      (NodeKey_Module . flip GWIB NotBoot     . unLoc <$> ms_home_imps ms)
-  where
-    -- Drop hs-boot nodes by using HsSrcFile as the key
-    hs_boot_key | drop_hs_boot_nodes = NotBoot -- is regular mod or signature
-                | otherwise          = IsBoot
-
-moduleGraphNodes :: Bool -> [ModuleGraphNode]
-  -> (Graph SummaryNode, NodeKey -> Maybe SummaryNode)
-moduleGraphNodes drop_hs_boot_nodes summaries =
-  (graphFromEdgedVerticesUniq nodes, lookup_node)
-  where
-    numbered_summaries = zip summaries [1..]
-
-    lookup_node :: NodeKey -> Maybe SummaryNode
-    lookup_node key = Map.lookup key (unNodeMap node_map)
-
-    lookup_key :: NodeKey -> Maybe Int
-    lookup_key = fmap summaryNodeKey . lookup_node
-
-    node_map :: NodeMap SummaryNode
-    node_map = NodeMap $
-      Map.fromList [ (mkNodeKey s, node)
-                   | node <- nodes
-                   , let s = summaryNodeSummary node
-                   ]
-
-    -- We use integers as the keys for the SCC algorithm
-    nodes :: [SummaryNode]
-    nodes = [ DigraphNode s key $ out_edge_keys $ unfilteredEdges drop_hs_boot_nodes s
-            | (s, key) <- numbered_summaries
-             -- Drop the hi-boot ones if told to do so
-            , case s of
-                InstantiationNode _ -> True
-                ModuleNode ems -> not $ isBootSummary (emsModSummary ems) == IsBoot && drop_hs_boot_nodes
-            ]
-
-    out_edge_keys :: [NodeKey] -> [Int]
-    out_edge_keys = mapMaybe lookup_key
-        -- If we want keep_hi_boot_nodes, then we do lookup_key with
-        -- IsBoot; else False
-
 -- The nodes of the graph are keyed by (mod, is boot?) pairs for the current
 -- modules, and indefinite unit IDs for dependencies which are instantiated with
 -- our holes.
 --
 -- NB: hsig files show up as *normal* nodes (not boot!), since they don't
 -- participate in cycles (for now)
-type ModNodeKey = ModuleNameWithIsBoot
+
+mkNodeMap :: [ExtendedModSummary] -> ModNodeMap ExtendedModSummary
+mkNodeMap summaries = ModNodeMap $ Map.fromList
+  [ (ms_mnwib $ emsModSummary s, s) | s <- summaries]
+
 newtype ModNodeMap a = ModNodeMap { unModNodeMap :: Map.Map ModNodeKey a }
   deriving (Functor, Traversable, Foldable)
 
@@ -1411,37 +1348,6 @@ modNodeMapElems (ModNodeMap m) = Map.elems m
 modNodeMapLookup :: ModNodeKey -> ModNodeMap a -> Maybe a
 modNodeMapLookup k (ModNodeMap m) = Map.lookup k m
 
-data NodeKey = NodeKey_Unit {-# UNPACK #-} !InstantiatedUnit | NodeKey_Module {-# UNPACK #-} !ModNodeKey
-  deriving (Eq, Ord)
-
-instance Outputable NodeKey where
-  ppr nk = pprNodeKey nk
-
-newtype NodeMap a = NodeMap { unNodeMap :: Map.Map NodeKey a }
-  deriving (Functor, Traversable, Foldable)
-
-mkNodeKey :: ModuleGraphNode -> NodeKey
-mkNodeKey = \case
-  InstantiationNode x -> NodeKey_Unit x
-  ModuleNode x -> NodeKey_Module $ mkHomeBuildModule0 (emsModSummary x)
-
-mkHomeBuildModule0 :: ModSummary -> ModuleNameWithIsBoot
-mkHomeBuildModule0 ms = GWIB
-  { gwib_mod = moduleName $ ms_mod ms
-  , gwib_isBoot = isBootSummary ms
-  }
-
-msKey :: ModSummary -> ModuleNameWithIsBoot
-msKey = mkHomeBuildModule0
-
-pprNodeKey :: NodeKey -> SDoc
-pprNodeKey (NodeKey_Unit iu) = ppr iu
-pprNodeKey (NodeKey_Module mk) = ppr mk
-
-mkNodeMap :: [ExtendedModSummary] -> ModNodeMap ExtendedModSummary
-mkNodeMap summaries = ModNodeMap $ Map.fromList
-  [ (mkHomeBuildModule0 $ emsModSummary s, s) | s <- summaries]
-
 -- | Efficiently construct a map from a NodeKey to its list of transitive dependencies
 mkDepsMap :: [ModuleGraphNode] -> (NodeKey -> [NodeKey])
 mkDepsMap nodes =
diff --git a/compiler/GHC/Unit/Module/Graph.hs b/compiler/GHC/Unit/Module/Graph.hs
index 822f72b88b..abee5d97aa 100644
--- a/compiler/GHC/Unit/Module/Graph.hs
+++ b/compiler/GHC/Unit/Module/Graph.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE LambdaCase      #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE DeriveTraversable #-}
 
 module GHC.Unit.Module.Graph
    ( ModuleGraph
@@ -18,10 +20,23 @@ module GHC.Unit.Module.Graph
    , mgElemModule
    , mgLookupModule
    , mgBootModules
+   , mgTransDeps
    , needsTemplateHaskellOrQQ
    , isTemplateHaskellOrQQNonBoot
    , showModMsg
-   , moduleGraphNodeModule)
+   , moduleGraphNodeModule
+
+   , moduleGraphNodes
+   , SummaryNode
+   , summaryNodeSummary
+
+   , NodeKey(..)
+   , ModNodeKey
+   , mkNodeKey
+   , msKey
+
+   )
+
 where
 
 import GHC.Prelude
@@ -29,13 +44,13 @@ import GHC.Prelude
 import qualified GHC.LanguageExtensions as LangExt
 
 import GHC.Data.Maybe
-import GHC.Data.Graph.Directed ( SCC(..) )
+import GHC.Data.Graph.Directed
 
 import GHC.Driver.Backend
 import GHC.Driver.Ppr
 import GHC.Driver.Session
 
-import GHC.Types.SourceFile ( hscSourceString )
+import GHC.Types.SourceFile ( hscSourceString, HscSource (HsBootFile) )
 
 import GHC.Unit.Module.ModSummary
 import GHC.Unit.Module.Env
@@ -43,6 +58,10 @@ import GHC.Unit.Types
 import GHC.Utils.Outputable
 
 import System.FilePath
+import qualified Data.Map as Map
+import GHC.Types.Unique.DSet
+import GHC.Types.SrcLoc
+import qualified Data.Set as Set
 
 -- | A '@ModuleGraphNode@' is a node in the '@ModuleGraph@'.
 -- Edges between nodes mark dependencies arising from module imports
@@ -76,6 +95,9 @@ instance Outputable ModuleGraphNode where
 -- 'GHC.topSortModuleGraph' and 'GHC.Data.Graph.Directed.flattenSCC' to achieve this.
 data ModuleGraph = ModuleGraph
   { mg_mss :: [ModuleGraphNode]
+  , mg_trans_deps :: Map.Map NodeKey (Set.Set NodeKey)
+    -- A cached transitive dependency calculation so that a lot of work is not
+    -- repeated whenever the transitive dependencies need to be calculated (for example, hptInstances)
   , mg_non_boot :: ModuleEnv ModSummary
     -- a map of all non-boot ModSummaries keyed by Modules
   , mg_boot :: ModuleSet
@@ -107,6 +129,9 @@ mapMG f mg@ModuleGraph{..} = mg
 mgBootModules :: ModuleGraph -> ModuleSet
 mgBootModules ModuleGraph{..} = mg_boot
 
+mgTransDeps :: ModuleGraph -> Map.Map NodeKey (Set.Set NodeKey)
+mgTransDeps = mg_trans_deps
+
 mgModSummaries :: ModuleGraph -> [ModSummary]
 mgModSummaries mg = [ m | ModuleNode (ExtendedModSummary m _) <- mgModSummaries' mg ]
 
@@ -124,7 +149,7 @@ mgLookupModule :: ModuleGraph -> Module -> Maybe ModSummary
 mgLookupModule ModuleGraph{..} m = lookupModuleEnv mg_non_boot m
 
 emptyMG :: ModuleGraph
-emptyMG = ModuleGraph [] emptyModuleEnv emptyModuleSet False
+emptyMG = ModuleGraph [] Map.empty emptyModuleEnv emptyModuleSet False
 
 isTemplateHaskellOrQQNonBoot :: ModSummary -> Bool
 isTemplateHaskellOrQQNonBoot ms =
@@ -137,6 +162,7 @@ isTemplateHaskellOrQQNonBoot ms =
 extendMG :: ModuleGraph -> ExtendedModSummary -> ModuleGraph
 extendMG ModuleGraph{..} ems@(ExtendedModSummary ms _) = ModuleGraph
   { mg_mss = ModuleNode ems : mg_mss
+  , mg_trans_deps = td
   , mg_non_boot = case isBootSummary ms of
       IsBoot -> mg_non_boot
       NotBoot -> extendModuleEnv mg_non_boot (ms_mod ms) ms
@@ -145,6 +171,9 @@ extendMG ModuleGraph{..} ems@(ExtendedModSummary ms _) = ModuleGraph
       IsBoot -> extendModuleSet mg_boot (ms_mod ms)
   , mg_needs_th_or_qq = mg_needs_th_or_qq || isTemplateHaskellOrQQNonBoot ms
   }
+  where
+    (gg, _lookup_node) = moduleGraphNodes False (ModuleNode ems : mg_mss)
+    td = allReachable gg (mkNodeKey . node_payload)
 
 extendMGInst :: ModuleGraph -> InstantiatedUnit -> ModuleGraph
 extendMGInst mg depUnitId = mg
@@ -209,3 +238,96 @@ showModMsg dflags recomp (ModuleNode (ExtendedModSummary mod_summary _)) =
                     then text obj_file <> comma <+> text dyn_file
                     else text obj_file
 
+
+
+
+type SummaryNode = Node Int ModuleGraphNode
+
+summaryNodeKey :: SummaryNode -> Int
+summaryNodeKey = node_key
+
+summaryNodeSummary :: SummaryNode -> ModuleGraphNode
+summaryNodeSummary = node_payload
+
+-- | Collect the immediate dependencies of a ModuleGraphNode,
+-- optionally avoiding hs-boot dependencies.
+-- If the drop_hs_boot_nodes flag is False, and if this is a .hs and there is
+-- an equivalent .hs-boot, add a link from the former to the latter.  This
+-- has the effect of detecting bogus cases where the .hs-boot depends on the
+-- .hs, by introducing a cycle.  Additionally, it ensures that we will always
+-- process the .hs-boot before the .hs, and so the HomePackageTable will always
+-- have the most up to date information.
+unfilteredEdges :: Bool -> ModuleGraphNode -> [NodeKey]
+unfilteredEdges drop_hs_boot_nodes = \case
+    InstantiationNode iuid ->
+      NodeKey_Module . flip GWIB NotBoot <$> uniqDSetToList (instUnitHoles iuid)
+    ModuleNode (ExtendedModSummary ms bds) ->
+      [ NodeKey_Unit inst_unit | inst_unit <- bds ] ++
+      (NodeKey_Module . flip GWIB hs_boot_key . unLoc <$> ms_home_srcimps ms) ++
+      [ NodeKey_Module $ GWIB (ms_mod_name ms) IsBoot
+      | not $ drop_hs_boot_nodes || ms_hsc_src ms == HsBootFile
+      ] ++
+      (NodeKey_Module . flip GWIB NotBoot     . unLoc <$> ms_home_imps ms)
+  where
+    -- Drop hs-boot nodes by using HsSrcFile as the key
+    hs_boot_key | drop_hs_boot_nodes = NotBoot -- is regular mod or signature
+                | otherwise          = IsBoot
+
+moduleGraphNodes :: Bool -> [ModuleGraphNode]
+  -> (Graph SummaryNode, NodeKey -> Maybe SummaryNode)
+moduleGraphNodes drop_hs_boot_nodes summaries =
+  (graphFromEdgedVerticesUniq nodes, lookup_node)
+  where
+    numbered_summaries = zip summaries [1..]
+
+    lookup_node :: NodeKey -> Maybe SummaryNode
+    lookup_node key = Map.lookup key (unNodeMap node_map)
+
+    lookup_key :: NodeKey -> Maybe Int
+    lookup_key = fmap summaryNodeKey . lookup_node
+
+    node_map :: NodeMap SummaryNode
+    node_map = NodeMap $
+      Map.fromList [ (mkNodeKey s, node)
+                   | node <- nodes
+                   , let s = summaryNodeSummary node
+                   ]
+
+    -- We use integers as the keys for the SCC algorithm
+    nodes :: [SummaryNode]
+    nodes = [ DigraphNode s key $ out_edge_keys $ unfilteredEdges drop_hs_boot_nodes s
+            | (s, key) <- numbered_summaries
+             -- Drop the hi-boot ones if told to do so
+            , case s of
+                InstantiationNode _ -> True
+                ModuleNode ems -> not $ isBootSummary (emsModSummary ems) == IsBoot && drop_hs_boot_nodes
+            ]
+
+    out_edge_keys :: [NodeKey] -> [Int]
+    out_edge_keys = mapMaybe lookup_key
+        -- If we want keep_hi_boot_nodes, then we do lookup_key with
+        -- IsBoot; else False
+
+type ModNodeKey = ModuleNameWithIsBoot
+
+data NodeKey = NodeKey_Unit {-# UNPACK #-} !InstantiatedUnit | NodeKey_Module {-# UNPACK #-} !ModNodeKey
+  deriving (Eq, Ord)
+
+instance Outputable NodeKey where
+  ppr nk = pprNodeKey nk
+
+newtype NodeMap a = NodeMap { unNodeMap :: Map.Map NodeKey a }
+  deriving (Functor, Traversable, Foldable)
+
+mkNodeKey :: ModuleGraphNode -> NodeKey
+mkNodeKey = \case
+  InstantiationNode x -> NodeKey_Unit x
+  ModuleNode x -> NodeKey_Module $ ms_mnwib (emsModSummary x)
+
+msKey :: ModSummary -> ModuleNameWithIsBoot
+msKey = ms_mnwib
+
+pprNodeKey :: NodeKey -> SDoc
+pprNodeKey (NodeKey_Unit iu) = ppr iu
+pprNodeKey (NodeKey_Module mk) = ppr mk
+