Add version of `reachableGraph` that avoids loop for cyclic inputs

by building its result connected component by component

Fixes #22512

1 files changed, 77 insertions, 26 deletions

diff --git a/compiler/GHC/Data/Graph/Directed.hs b/compiler/GHC/Data/Graph/Directed.hs
index 85685691c3..1f4202038e 100644
--- a/compiler/GHC/Data/Graph/Directed.hs
+++ b/compiler/GHC/Data/Graph/Directed.hs
@@ -9,11 +9,11 @@ module GHC.Data.Graph.Directed (
         Graph, graphFromEdgedVerticesOrd, graphFromEdgedVerticesUniq,
         graphFromVerticesAndAdjacency,
 
-        SCC(..), Node(..), flattenSCC, flattenSCCs,
+        SCC(..), Node(..), G.flattenSCC, G.flattenSCCs,
         stronglyConnCompG,
         topologicalSortG,
         verticesG, edgesG, hasVertexG,
-        reachableG, reachablesG, transposeG, allReachable, outgoingG,
+        reachableG, reachablesG, transposeG, allReachable, allReachableCyclic, outgoingG,
         emptyG,
 
         findCycle,
@@ -58,7 +58,7 @@ import qualified Data.Map as Map
 import qualified Data.Set as Set
 
 import qualified Data.Graph as G
-import Data.Graph hiding (Graph, Edge, transposeG, reachable)
+import Data.Graph ( Vertex, Bounds, SCC(..) ) -- Used in the underlying representation
 import Data.Tree
 import GHC.Types.Unique
 import GHC.Types.Unique.FM
@@ -291,19 +291,11 @@ We use the order of nodes to normalize the order of edges.
 -}
 
 stronglyConnCompG :: Graph node -> [SCC node]
-stronglyConnCompG graph = decodeSccs graph forest
-  where forest = {-# SCC "Digraph.scc" #-} scc (gr_int_graph graph)
-
-decodeSccs :: Graph node -> Forest Vertex -> [SCC node]
-decodeSccs Graph { gr_int_graph = graph, gr_vertex_to_node = vertex_fn } forest
-  = map decode forest
-  where
-    decode (Node v []) | mentions_itself v = CyclicSCC [vertex_fn v]
-                       | otherwise         = AcyclicSCC (vertex_fn v)
-    decode other = CyclicSCC (dec other [])
-      where dec (Node v ts) vs = vertex_fn v : foldr dec vs ts
-    mentions_itself v = v `elem` (graph ! v)
+stronglyConnCompG graph = decodeSccs graph $ scc (gr_int_graph graph)
 
+decodeSccs :: Graph node -> [SCC Vertex] -> [SCC node]
+decodeSccs Graph { gr_vertex_to_node = vertex_fn }
+  = map (fmap vertex_fn)
 
 -- The following two versions are provided for backwards compatibility:
 -- See Note [Deterministic SCC]
@@ -357,7 +349,7 @@ stronglyConnCompFromEdgedVerticesUniqR =
 
 topologicalSortG :: Graph node -> [node]
 topologicalSortG graph = map (gr_vertex_to_node graph) result
-  where result = {-# SCC "Digraph.topSort" #-} topSort (gr_int_graph graph)
+  where result = {-# SCC "Digraph.topSort" #-} G.topSort (gr_int_graph graph)
 
 reachableG :: Graph node -> node -> [node]
 reachableG graph from = map (gr_vertex_to_node graph) result
@@ -377,22 +369,31 @@ reachablesG graph froms = map (gr_vertex_to_node graph) result
         vs = [ v | Just v <- map (gr_node_to_vertex graph) froms ]
 
 -- | Efficiently construct a map which maps each key to it's set of transitive
--- dependencies.
+-- dependencies. Only works on acyclic input.
 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]
+allReachable = all_reachable reachableGraph
+
+-- | Efficiently construct a map which maps each key to it's set of transitive
+-- dependencies. Less efficient than @allReachable@, but works on cyclic input as well.
+allReachableCyclic :: Ord key => Graph node -> (node -> key) -> M.Map key (S.Set key)
+allReachableCyclic = all_reachable reachableGraphCyclic
+
+all_reachable :: Ord key => (IntGraph -> IM.IntMap IS.IntSet) -> Graph node -> (node -> key) -> M.Map key (S.Set key)
+all_reachable int_reachables (Graph g from _) keyOf =
+  M.fromList [(k, IS.foldr (\v' vs -> keyOf (from v') `S.insert` vs) S.empty vs)
+             | (v, vs) <- IM.toList int_graph
+             , let k = keyOf (from v)]
   where
-    int_graph = reachableGraph g
+    int_graph = int_reachables g
 
 hasVertexG :: Graph node -> node -> Bool
 hasVertexG graph node = isJust $ gr_node_to_vertex graph node
 
 verticesG :: Graph node -> [node]
-verticesG graph = map (gr_vertex_to_node graph) $ vertices (gr_int_graph graph)
+verticesG graph = map (gr_vertex_to_node graph) $ G.vertices (gr_int_graph graph)
 
 edgesG :: Graph node -> [Edge node]
-edgesG graph = map (\(v1, v2) -> Edge (v2n v1) (v2n v2)) $ edges (gr_int_graph graph)
+edgesG graph = map (\(v1, v2) -> Edge (v2n v1) (v2n v2)) $ G.edges (gr_int_graph graph)
   where v2n = gr_vertex_to_node graph
 
 transposeG :: Graph node -> Graph node
@@ -452,13 +453,63 @@ preorderF ts         = concatMap flatten ts
 
 -- This generalizes reachable which was found in Data.Graph
 reachable    :: IntGraph -> [Vertex] -> [Vertex]
-reachable g vs = preorderF (dfs g vs)
+reachable g vs = preorderF (G.dfs g vs)
 
 reachableGraph :: IntGraph -> IM.IntMap IS.IntSet
 reachableGraph g = res
   where
     do_one v = IS.unions (IS.fromList (g ! v) : mapMaybe (flip IM.lookup res) (g ! v))
-    res = IM.fromList [(v, do_one v) | v <- vertices g]
+    res = IM.fromList [(v, do_one v) | v <- G.vertices g]
+
+scc :: IntGraph -> [SCC Vertex]
+scc graph = map decode forest
+  where
+    forest = {-# SCC "Digraph.scc" #-} G.scc graph
+
+    decode (Node v []) | mentions_itself v = CyclicSCC [v]
+                       | otherwise         = AcyclicSCC v
+    decode other = CyclicSCC (dec other [])
+      where dec (Node v ts) vs = v : foldr dec vs ts
+    mentions_itself v = v `elem` (graph ! v)
+
+reachableGraphCyclic :: IntGraph -> IM.IntMap IS.IntSet
+reachableGraphCyclic g = foldl' add_one_comp mempty comps
+  where
+    neighboursOf v = g!v
+
+    comps = scc g
+
+    -- To avoid divergence on cyclic input, we build the result
+    -- strongly connected component by component, in topological
+    -- order. For each SCC, we know that:
+    --
+    --   * All vertices in the component can reach all other vertices
+    --     in the component ("local" reachables)
+    --
+    --   * Other reachable vertices ("remote" reachables) must come
+    --     from earlier components, either via direct neighbourhood, or
+    --     transitively from earlier reachability map
+    --
+    -- This allows us to build the extension of the reachability map
+    -- directly, without any self-reference, thereby avoiding a loop.
+    add_one_comp :: IM.IntMap IS.IntSet -> SCC Vertex -> IM.IntMap IS.IntSet
+    add_one_comp earlier (AcyclicSCC v) = IM.insert v all_remotes earlier
+      where
+        earlier_neighbours = neighboursOf v
+        earlier_further = mapMaybe (flip IM.lookup earlier) earlier_neighbours
+        all_remotes = IS.unions (IS.fromList earlier_neighbours : earlier_further)
+    add_one_comp earlier (CyclicSCC vs) = IM.union (IM.fromList [(v, local v `IS.union` all_remotes) | v <- vs]) earlier
+      where
+        all_locals = IS.fromList vs
+        local v = IS.delete v all_locals
+            -- Arguably, for a cyclic SCC we should include each
+            -- vertex in its own reachable set. However, this could
+            -- lead to a lot of extra pain in client code to avoid
+            -- looping when traversing the reachability map.
+        all_neighbours = IS.fromList (concatMap neighboursOf vs)
+        earlier_neighbours = all_neighbours IS.\\ all_locals
+        earlier_further = mapMaybe (flip IM.lookup earlier) (IS.toList earlier_neighbours)
+        all_remotes = IS.unions (earlier_neighbours : earlier_further)
 
 {-
 ************************************************************************
@@ -565,4 +616,4 @@ graphFromVerticesAndAdjacency vertices edges = Graph graph vertex_node (key_vert
         key_vertex_pair (a, b) = (expectJust "graphFromVerticesAndAdjacency" $ key_vertex a,
                                   expectJust "graphFromVerticesAndAdjacency" $ key_vertex b)
         reduced_edges = map key_vertex_pair edges
-        graph = buildG bounds reduced_edges
+        graph = G.buildG bounds reduced_edges