{-# LANGUAGE OverloadedStrings #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

-- |
-- Module      :  Disco.Typecheck.Graph
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- A thin layer on top of graphs from the @fgl@ package, which
--   allows dealing with vertices by label instead of by integer
--   @Node@ values.
module Disco.Typecheck.Graph where

import Prelude hiding (map, (<>))
import qualified Prelude as P

import Control.Arrow ((&&&))
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe (fromJust, isJust, mapMaybe)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Tuple (swap)

import qualified Data.Graph.Inductive.Graph as G
import Data.Graph.Inductive.PatriciaTree (Gr)
import qualified Data.Graph.Inductive.Query.DFS as G (
  components,
  condensation,
  topsort',
 )

import Disco.Pretty
import Disco.Util ((!))

-- | Directed graphs, with vertices labelled by @a@ and unlabelled
--   edges.
data Graph a = G (Gr a ()) (Map a G.Node) (Map G.Node a)
  deriving (Int -> Graph a -> ShowS
[Graph a] -> ShowS
Graph a -> String
(Int -> Graph a -> ShowS)
-> (Graph a -> String) -> ([Graph a] -> ShowS) -> Show (Graph a)
forall a. Show a => Int -> Graph a -> ShowS
forall a. Show a => [Graph a] -> ShowS
forall a. Show a => Graph a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Graph a -> ShowS
showsPrec :: Int -> Graph a -> ShowS
$cshow :: forall a. Show a => Graph a -> String
show :: Graph a -> String
$cshowList :: forall a. Show a => [Graph a] -> ShowS
showList :: [Graph a] -> ShowS
Show)

instance Pretty a => Pretty (Graph a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Graph a -> Sem r (Doc ann)
pretty (G Gr a ()
g Map a Int
_ Map Int a
_) = Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
parens (Sem r (Doc ann)
prettyVertices Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann)
", " Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann)
prettyEdges)
   where
    -- (V = {(0, x), (1, N)}, E = {0 -> 1, 2 -> 3})

    vs :: [LNode a]
vs = Gr a () -> [LNode a]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes Gr a ()
g
    es :: [LEdge ()]
es = Gr a () -> [LEdge ()]
forall a b. Gr a b -> [LEdge b]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LEdge b]
G.labEdges Gr a ()
g

    prettyVertex :: (a, t) -> Sem r (Doc ann)
prettyVertex (a
n, t
a) = Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
parens (String -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => String -> m (Doc ann)
text (a -> String
forall a. Show a => a -> String
show a
n) Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann)
", " Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> t -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
t -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty t
a)
    prettyVertices :: Sem r (Doc ann)
prettyVertices = Sem r (Doc ann)
"V = " Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
braces (Sem r (Doc ann) -> [Sem r (Doc ann)] -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Monad f =>
f (Doc ann) -> [f (Doc ann)] -> f (Doc ann)
intercalate Sem r (Doc ann)
"," ((LNode a -> Sem r (Doc ann)) -> [LNode a] -> [Sem r (Doc ann)]
forall a b. (a -> b) -> [a] -> [b]
P.map LNode a -> Sem r (Doc ann)
forall {a} {t} {r :: EffectRow} {ann}.
(Show a, Pretty t, Member (Reader PA) r, Member LFresh r) =>
(a, t) -> Sem r (Doc ann)
prettyVertex [LNode a]
vs))
    prettyEdge :: (a, a, c) -> f (Doc ann)
prettyEdge (a
v1, a
v2, c
_) = String -> f (Doc ann)
forall (m :: * -> *) ann. Applicative m => String -> m (Doc ann)
text (a -> String
forall a. Show a => a -> String
show a
v1) f (Doc ann) -> f (Doc ann) -> f (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> f (Doc ann)
"->" f (Doc ann) -> f (Doc ann) -> f (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> String -> f (Doc ann)
forall (m :: * -> *) ann. Applicative m => String -> m (Doc ann)
text (a -> String
forall a. Show a => a -> String
show a
v2)
    prettyEdges :: Sem r (Doc ann)
prettyEdges = Sem r (Doc ann)
"E = " Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
braces (Sem r (Doc ann) -> [Sem r (Doc ann)] -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Monad f =>
f (Doc ann) -> [f (Doc ann)] -> f (Doc ann)
intercalate Sem r (Doc ann)
"," ((LEdge () -> Sem r (Doc ann)) -> [LEdge ()] -> [Sem r (Doc ann)]
forall a b. (a -> b) -> [a] -> [b]
P.map LEdge () -> Sem r (Doc ann)
forall {f :: * -> *} {ann} {a} {a} {c}.
(Applicative f, IsString (f (Doc ann)), Show a, Show a) =>
(a, a, c) -> f (Doc ann)
prettyEdge [LEdge ()]
es))

-- | Create a graph with the given set of vertices and directed edges.
--   If any edges refer to vertices that are not in the given vertex
--   set, they will simply be dropped.
mkGraph :: (Show a, Ord a) => Set a -> Set (a, a) -> Graph a
mkGraph :: forall a. (Show a, Ord a) => Set a -> Set (a, a) -> Graph a
mkGraph Set a
vs Set (a, a)
es = Gr a () -> Map a Int -> Map Int a -> Graph a
forall a. Gr a () -> Map a Int -> Map Int a -> Graph a
G ([LNode a] -> [LEdge ()] -> Gr a ()
forall a b. [LNode a] -> [LEdge b] -> Gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
[LNode a] -> [LEdge b] -> gr a b
G.mkGraph [LNode a]
vs' [LEdge ()]
es') Map a Int
a2n Map Int a
n2a
 where
  vs' :: [LNode a]
vs' = [Int] -> [a] -> [LNode a]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
vs)
  n2a :: Map Int a
n2a = [LNode a] -> Map Int a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [LNode a]
vs'
  a2n :: Map a Int
a2n = [(a, Int)] -> Map a Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, Int)] -> Map a Int)
-> ([LNode a] -> [(a, Int)]) -> [LNode a] -> Map a Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LNode a -> (a, Int)) -> [LNode a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
P.map LNode a -> (a, Int)
forall a b. (a, b) -> (b, a)
swap ([LNode a] -> Map a Int) -> [LNode a] -> Map a Int
forall a b. (a -> b) -> a -> b
$ [LNode a]
vs'
  es' :: [LEdge ()]
es' = ((a, a) -> Maybe (LEdge ())) -> [(a, a)] -> [LEdge ()]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (a, a) -> Maybe (LEdge ())
mkEdge (Set (a, a) -> [(a, a)]
forall a. Set a -> [a]
S.toList Set (a, a)
es)
  mkEdge :: (a, a) -> Maybe (LEdge ())
mkEdge (a
a1, a
a2) = (,,) (Int -> Int -> () -> LEdge ())
-> Maybe Int -> Maybe (Int -> () -> LEdge ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Map a Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a1 Map a Int
a2n Maybe (Int -> () -> LEdge ())
-> Maybe Int -> Maybe (() -> LEdge ())
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> Map a Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a2 Map a Int
a2n Maybe (() -> LEdge ()) -> Maybe () -> Maybe (LEdge ())
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> () -> Maybe ()
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Return the set of vertices (nodes) of a graph.
nodes :: Graph a -> Set a
nodes :: forall a. Graph a -> Set a
nodes (G Gr a ()
_ Map a Int
m Map Int a
_) = Map a Int -> Set a
forall k a. Map k a -> Set k
M.keysSet Map a Int
m

-- | Return the set of directed edges of a graph.
edges :: Ord a => Graph a -> Set (a, a)
edges :: forall a. Ord a => Graph a -> Set (a, a)
edges (G Gr a ()
g Map a Int
_ Map Int a
m) = [(a, a)] -> Set (a, a)
forall a. Ord a => [a] -> Set a
S.fromList ([(a, a)] -> Set (a, a)) -> [(a, a)] -> Set (a, a)
forall a b. (a -> b) -> a -> b
$ (LEdge () -> (a, a)) -> [LEdge ()] -> [(a, a)]
forall a b. (a -> b) -> [a] -> [b]
P.map (\(Int
n1, Int
n2, ()) -> (Map Int a
m Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
! Int
n1, Map Int a
m Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
! Int
n2)) (Gr a () -> [LEdge ()]
forall a b. Gr a b -> [LEdge b]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LEdge b]
G.labEdges Gr a ()
g)

-- | Map a function over all the vertices of a graph.  @Graph@ is not
--   a @Functor@ instance because of the @Ord@ constraint on @b@.
map :: Ord b => (a -> b) -> Graph a -> Graph b
map :: forall b a. Ord b => (a -> b) -> Graph a -> Graph b
map a -> b
f (G Gr a ()
g Map a Int
m1 Map Int a
m2) = Gr b () -> Map b Int -> Map Int b -> Graph b
forall a. Gr a () -> Map a Int -> Map Int a -> Graph a
G ((a -> b) -> Gr a () -> Gr b ()
forall (gr :: * -> * -> *) a c b.
DynGraph gr =>
(a -> c) -> gr a b -> gr c b
G.nmap a -> b
f Gr a ()
g) ((a -> b) -> Map a Int -> Map b Int
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys a -> b
f Map a Int
m1) ((a -> b) -> Map Int a -> Map Int b
forall a b k. (a -> b) -> Map k a -> Map k b
M.map a -> b
f Map Int a
m2)

-- | Delete a vertex.
delete :: (Show a, Ord a) => a -> Graph a -> Graph a
delete :: forall a. (Show a, Ord a) => a -> Graph a -> Graph a
delete a
a (G Gr a ()
g Map a Int
a2n Map Int a
n2a) = Gr a () -> Map a Int -> Map Int a -> Graph a
forall a. Gr a () -> Map a Int -> Map Int a -> Graph a
G (Int -> Gr a () -> Gr a ()
forall (gr :: * -> * -> *) a b. Graph gr => Int -> gr a b -> gr a b
G.delNode Int
n Gr a ()
g) (a -> Map a Int -> Map a Int
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
a Map a Int
a2n) (Int -> Map Int a -> Map Int a
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Int
n Map Int a
n2a)
 where
  n :: Int
n = Map a Int
a2n Map a Int -> a -> Int
forall k v. (Show k, Ord k) => Map k v -> k -> v
! a
a

-- | The @condensation@ of a graph is the graph of its strongly
--   connected components, /i.e./ each strongly connected component is
--   compressed to a single node, labelled by the set of vertices in
--   the component.  There is an edge from component A to component B
--   in the condensed graph iff there is an edge from any vertex in
--   component A to any vertex in component B in the original graph.
condensation :: Ord a => Graph a -> Graph (Set a)
condensation :: forall a. Ord a => Graph a -> Graph (Set a)
condensation (G Gr a ()
g Map a Int
_ Map Int a
n2a) = Gr (Set a) ()
-> Map (Set a) Int -> Map Int (Set a) -> Graph (Set a)
forall a. Gr a () -> Map a Int -> Map Int a -> Graph a
G Gr (Set a) ()
g' Map (Set a) Int
as2n Map Int (Set a)
n2as
 where
  g' :: Gr (Set a) ()
g' = ([Int] -> Set a) -> Gr [Int] () -> Gr (Set a) ()
forall (gr :: * -> * -> *) a c b.
DynGraph gr =>
(a -> c) -> gr a b -> gr c b
G.nmap ([a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> ([Int] -> [a]) -> [Int] -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
P.map (Map Int a
n2a Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!)) (Gr a () -> Gr [Int] ()
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> gr [Int] ()
G.condensation Gr a ()
g)
  vs' :: [LNode (Set a)]
vs' = Gr (Set a) () -> [LNode (Set a)]
forall a b. Gr a b -> [LNode a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [LNode a]
G.labNodes Gr (Set a) ()
g'
  n2as :: Map Int (Set a)
n2as = [LNode (Set a)] -> Map Int (Set a)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [LNode (Set a)]
vs'
  as2n :: Map (Set a) Int
as2n = [(Set a, Int)] -> Map (Set a) Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Set a, Int)] -> Map (Set a) Int)
-> ([LNode (Set a)] -> [(Set a, Int)])
-> [LNode (Set a)]
-> Map (Set a) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LNode (Set a) -> (Set a, Int))
-> [LNode (Set a)] -> [(Set a, Int)]
forall a b. (a -> b) -> [a] -> [b]
P.map LNode (Set a) -> (Set a, Int)
forall a b. (a, b) -> (b, a)
swap ([LNode (Set a)] -> Map (Set a) Int)
-> [LNode (Set a)] -> Map (Set a) Int
forall a b. (a -> b) -> a -> b
$ [LNode (Set a)]
vs'

-- | Get a list of the weakly connected components of a graph,
--   providing the set of vertices in each.  Equivalently, return the
--   strongly connected components of the graph when considered as an
--   undirected graph.
wcc :: Ord a => Graph a -> [Set a]
wcc :: forall a. Ord a => Graph a -> [Set a]
wcc = (Set (Int, a) -> Set a) -> [Set (Int, a)] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
P.map (((Int, a) -> a) -> Set (Int, a) -> Set a
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (Int, a) -> a
forall a b. (a, b) -> b
snd) ([Set (Int, a)] -> [Set a])
-> (Graph a -> [Set (Int, a)]) -> Graph a -> [Set a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph a -> [Set (Int, a)]
forall a. Ord a => Graph a -> [Set (Int, a)]
wccIDs

wccIDs :: Ord a => Graph a -> [Set (G.Node, a)]
wccIDs :: forall a. Ord a => Graph a -> [Set (Int, a)]
wccIDs (G Gr a ()
g Map a Int
_a2n Map Int a
n2a) = ([Int] -> Set (Int, a)) -> [[Int]] -> [Set (Int, a)]
forall a b. (a -> b) -> [a] -> [b]
P.map ([(Int, a)] -> Set (Int, a)
forall a. Ord a => [a] -> Set a
S.fromList ([(Int, a)] -> Set (Int, a))
-> ([Int] -> [(Int, a)]) -> [Int] -> Set (Int, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> (Int, a)) -> [Int] -> [(Int, a)]
forall a b. (a -> b) -> [a] -> [b]
P.map (Int -> Int
forall a. a -> a
id (Int -> Int) -> (Int -> a) -> Int -> (Int, a)
forall b c c'. (b -> c) -> (b -> c') -> b -> (c, c')
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& (Map Int a
n2a Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!))) (Gr a () -> [[Int]]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [[Int]]
G.components Gr a ()
g)

-- | Do a topological sort on a DAG.
topsort :: Graph a -> [a]
topsort :: forall a. Graph a -> [a]
topsort (G Gr a ()
g Map a Int
_a2n Map Int a
_n2a) = Gr a () -> [a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [a]
G.topsort' Gr a ()
g

-- | A miscellaneous utility function to turn a @Graph Maybe@ into a
--   @Maybe Graph@: the result is @Just@ iff all the vertices in the
--   input graph are.
sequenceGraph :: Ord a => Graph (Maybe a) -> Maybe (Graph a)
sequenceGraph :: forall a. Ord a => Graph (Maybe a) -> Maybe (Graph a)
sequenceGraph Graph (Maybe a)
g = case (Maybe a -> Bool) -> Set (Maybe a) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (Graph (Maybe a) -> Set (Maybe a)
forall a. Graph a -> Set a
nodes Graph (Maybe a)
g) of
  Bool
False -> Maybe (Graph a)
forall a. Maybe a
Nothing
  Bool
True -> Graph a -> Maybe (Graph a)
forall a. a -> Maybe a
Just (Graph a -> Maybe (Graph a)) -> Graph a -> Maybe (Graph a)
forall a b. (a -> b) -> a -> b
$ (Maybe a -> a) -> Graph (Maybe a) -> Graph a
forall b a. Ord b => (a -> b) -> Graph a -> Graph b
map Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust Graph (Maybe a)
g

-- | Get a list of all the /successors/ of a given node in the graph,
--   /i.e./ all the nodes reachable from the given node by a directed
--   path.  Does not include the given node itself.
suc :: (Show a, Ord a) => Graph a -> a -> [a]
suc :: forall a. (Show a, Ord a) => Graph a -> a -> [a]
suc (G Gr a ()
g Map a Int
a2n Map Int a
n2a) = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
P.map (Map Int a
n2a Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!) ([Int] -> [a]) -> (a -> [Int]) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gr a () -> Int -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int -> [Int]
G.suc Gr a ()
g (Int -> [Int]) -> (a -> Int) -> a -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Int
a2n Map a Int -> a -> Int
forall k v. (Show k, Ord k) => Map k v -> k -> v
!)

-- | Get a list of all the /predecessors/ of a given node in the
--   graph, /i.e./ all the nodes from which from the given node is
--   reachable by a directed path.  Does not include the given node
--   itself.
pre :: (Show a, Ord a) => Graph a -> a -> [a]
pre :: forall a. (Show a, Ord a) => Graph a -> a -> [a]
pre (G Gr a ()
g Map a Int
a2n Map Int a
n2a) = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
P.map (Map Int a
n2a Map Int a -> Int -> a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!) ([Int] -> [a]) -> (a -> [Int]) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gr a () -> Int -> [Int]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> Int -> [Int]
G.pre Gr a ()
g (Int -> [Int]) -> (a -> Int) -> a -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map a Int
a2n Map a Int -> a -> Int
forall k v. (Show k, Ord k) => Map k v -> k -> v
!)

-- | Given a graph, return two mappings: the first maps each vertex to
--   its set of successors; the second maps each vertex to its set of
--   predecessors.  Equivalent to
--
--   > (M.fromList *** M.fromList) . unzip . map (\a -> ((a, suc g a), (a, pre g a))) . nodes $ g
--
--   but much more efficient.
cessors :: (Show a, Ord a) => Graph a -> (Map a (Set a), Map a (Set a))
cessors :: forall a.
(Show a, Ord a) =>
Graph a -> (Map a (Set a), Map a (Set a))
cessors g :: Graph a
g@(G Gr a ()
gg Map a Int
_ Map Int a
_) = (Map a (Set a)
succs, Map a (Set a)
preds)
 where
  as :: [a]
as = Gr a () -> [a]
forall (gr :: * -> * -> *) a b. Graph gr => gr a b -> [a]
G.topsort' Gr a ()
gg
  succs :: Map a (Set a)
succs = (a -> Map a (Set a) -> Map a (Set a))
-> Map a (Set a) -> [a] -> Map a (Set a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> Map a (Set a) -> Map a (Set a)
collectSuccs Map a (Set a)
forall k a. Map k a
M.empty [a]
as -- build successors map
  collectSuccs :: a -> Map a (Set a) -> Map a (Set a)
collectSuccs a
a Map a (Set a)
m = a -> Set a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
a Set a
succsSet Map a (Set a)
m
   where
    ss :: [a]
ss = Graph a -> a -> [a]
forall a. (Show a, Ord a) => Graph a -> a -> [a]
suc Graph a
g a
a
    succsSet :: Set a
succsSet = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
ss Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [Set a] -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((a -> Set a) -> [a] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
P.map (Map a (Set a)
m Map a (Set a) -> a -> Set a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!) [a]
ss)

  preds :: Map a (Set a)
preds = (a -> Map a (Set a) -> Map a (Set a))
-> Map a (Set a) -> [a] -> Map a (Set a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> Map a (Set a) -> Map a (Set a)
collectPreds Map a (Set a)
forall k a. Map k a
M.empty ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
as) -- build predecessors map
  collectPreds :: a -> Map a (Set a) -> Map a (Set a)
collectPreds a
a Map a (Set a)
m = a -> Set a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
a Set a
predsSet Map a (Set a)
m
   where
    ss :: [a]
ss = Graph a -> a -> [a]
forall a. (Show a, Ord a) => Graph a -> a -> [a]
pre Graph a
g a
a
    predsSet :: Set a
predsSet = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
ss Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [Set a] -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((a -> Set a) -> [a] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
P.map (Map a (Set a)
m Map a (Set a) -> a -> Set a
forall k v. (Show k, Ord k) => Map k v -> k -> v
!) [a]
ss)