content: {-# 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 (Show)
  
  instance Pretty a => Pretty (Graph a) where
    pretty (G g _ _) = parens (prettyVertices <> ", " <> prettyEdges)
     where
      -- (V = {(0, x), (1, N)}, E = {0 -> 1, 2 -> 3})
  
      vs = G.labNodes g
      es = G.labEdges g
  
      prettyVertex (n, a) = parens (text (show n) <> ", " <> pretty a)
      prettyVertices = "V = " <> braces (intercalate "," (P.map prettyVertex vs))
      prettyEdge (v1, v2, _) = text (show v1) <+> "->" <+> text (show v2)
      prettyEdges = "E = " <> braces (intercalate "," (P.map prettyEdge 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 vs es = G (G.mkGraph vs' es') a2n n2a
   where
    vs' = zip [0 ..] (S.toList vs)
    n2a = M.fromList vs'
    a2n = M.fromList . P.map swap $ vs'
    es' = mapMaybe mkEdge (S.toList es)
    mkEdge (a1, a2) = (,,) <$> M.lookup a1 a2n <*> M.lookup a2 a2n <*> pure ()
  
  -- | Return the set of vertices (nodes) of a graph.
  nodes :: Graph a -> Set a
  nodes (G _ m _) = M.keysSet m
  
  -- | Return the set of directed edges of a graph.
  edges :: Ord a => Graph a -> Set (a, a)
  edges (G g _ m) = S.fromList $ P.map (\(n1, n2, ()) -> (m ! n1, m ! n2)) (G.labEdges 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 f (G g m1 m2) = G (G.nmap f g) (M.mapKeys f m1) (M.map f m2)
  
  -- | Delete a vertex.
  delete :: (Show a, Ord a) => a -> Graph a -> Graph a
  delete a (G g a2n n2a) = G (G.delNode n g) (M.delete a a2n) (M.delete n n2a)
   where
    n = a2n ! 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 (G g _ n2a) = G g' as2n n2as
   where
    g' = G.nmap (S.fromList . P.map (n2a !)) (G.condensation g)
    vs' = G.labNodes g'
    n2as = M.fromList vs'
    as2n = M.fromList . P.map swap $ 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 = P.map (S.map snd) . wccIDs
  
  wccIDs :: Ord a => Graph a -> [Set (G.Node, a)]
  wccIDs (G g _a2n n2a) = P.map (S.fromList . P.map (id &&& (n2a !))) (G.components g)
  
  -- | Do a topological sort on a DAG.
  topsort :: Graph a -> [a]
  topsort (G g _a2n _n2a) = G.topsort' 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 g = case all isJust (nodes g) of
    False -> Nothing
    True -> Just $ map fromJust 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 (G g a2n n2a) = P.map (n2a !) . G.suc g . (a2n !)
  
  -- | 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 (G g a2n n2a) = P.map (n2a !) . G.pre g . (a2n !)
  
  -- | 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 g@(G gg _ _) = (succs, preds)
   where
    as = G.topsort' gg
    succs = foldr collectSuccs M.empty as -- build successors map
    collectSuccs a m = M.insert a succsSet m
     where
      ss = suc g a
      succsSet = S.fromList ss `S.union` S.unions (P.map (m !) ss)
  
    preds = foldr collectPreds M.empty (reverse as) -- build predecessors map
    collectPreds a m = M.insert a predsSet m
     where
      ss = pre g a
      predsSet = S.fromList ss `S.union` S.unions (P.map (m !) ss)
  