-----------------------------------------------------------------------------
-- |
-- Module     : Data.Graph.Typed
-- Copyright  : (c) Anton Lorenzen, Andrey Mokhov 2016-2022
-- License    : MIT (see the file LICENSE)
-- Maintainer : anfelor@posteo.de, andrey.mokhov@gmail.com
-- Stability  : unstable
--
-- __Alga__ is a library for algebraic construction and manipulation of graphs
-- in Haskell. See <https://github.com/snowleopard/alga-paper this paper> for the
-- motivation behind the library, the underlying theory, and implementation details.
--
-- This module provides primitives for interoperability between this library and
-- the "Data.Graph" module of the containers library. It is for internal use only
-- and may be removed without notice at any point.
-----------------------------------------------------------------------------
module Data.Graph.Typed (
    -- * Data type and construction
    GraphKL(..), fromAdjacencyMap, fromAdjacencyIntMap,

    -- * Basic algorithms
    dfsForest, dfsForestFrom, dfs, topSort, scc
    ) where

import Data.Tree
import Data.Maybe
import Data.Foldable

import qualified Data.Graph as KL

import qualified Algebra.Graph.AdjacencyMap          as AM
import qualified Algebra.Graph.NonEmpty.AdjacencyMap as NonEmpty
import qualified Algebra.Graph.AdjacencyIntMap       as AIM

import qualified Data.Map.Strict                     as Map
import qualified Data.Set                            as Set

-- | 'GraphKL' encapsulates King-Launchbury graphs, which are implemented in
-- the "Data.Graph" module of the @containers@ library.
data GraphKL a = GraphKL {
    -- | Array-based graph representation (King and Launchbury, 1995).
    forall a. GraphKL a -> Graph
toGraphKL :: KL.Graph,
    -- | A mapping of "Data.Graph.Vertex" to vertices of type @a@.
    -- This is partial and may fail if the vertex is out of bounds.
    forall a. GraphKL a -> Int -> a
fromVertexKL :: KL.Vertex -> a,
    -- | A mapping from vertices of type @a@ to "Data.Graph.Vertex".
    -- Returns 'Nothing' if the argument is not in the graph.
    forall a. GraphKL a -> a -> Maybe Int
toVertexKL :: a -> Maybe KL.Vertex }

-- | Build 'GraphKL' from an 'AM.AdjacencyMap'. If @fromAdjacencyMap g == h@
-- then the following holds:
--
-- @
-- map ('fromVertexKL' h) ('Data.Graph.vertices' $ 'toGraphKL' h)                               == 'AM.vertexList' g
-- map (\\(x, y) -> ('fromVertexKL' h x, 'fromVertexKL' h y)) ('Data.Graph.edges' $ 'toGraphKL' h) == 'AM.edgeList' g
-- 'toGraphKL' (fromAdjacencyMap (1 * 2 + 3 * 1))                                == 'array' (0,2) [(0,[1]), (1,[]), (2,[0])]
-- 'toGraphKL' (fromAdjacencyMap (1 * 2 + 2 * 1))                                == 'array' (0,1) [(0,[1]), (1,[0])]
-- @
fromAdjacencyMap :: Ord a => AM.AdjacencyMap a -> GraphKL a
fromAdjacencyMap :: forall a. Ord a => AdjacencyMap a -> GraphKL a
fromAdjacencyMap AdjacencyMap a
am = GraphKL
    { toGraphKL :: Graph
toGraphKL    = Graph
g
    , fromVertexKL :: Int -> a
fromVertexKL = \Int
u -> case Int -> ((), a, [a])
r Int
u of (()
_, a
v, [a]
_) -> a
v
    , toVertexKL :: a -> Maybe Int
toVertexKL   = a -> Maybe Int
t }
  where
    (Graph
g, Int -> ((), a, [a])
r, a -> Maybe Int
t) = [((), a, [a])] -> (Graph, Int -> ((), a, [a]), a -> Maybe Int)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Int -> (node, key, [key]), key -> Maybe Int)
KL.graphFromEdges [ ((), a
x, [a]
ys) | (a
x, [a]
ys) <- AdjacencyMap a -> [(a, [a])]
forall a. AdjacencyMap a -> [(a, [a])]
AM.adjacencyList AdjacencyMap a
am ]

-- | Build 'GraphKL' from an 'AIM.AdjacencyIntMap'. If
-- @fromAdjacencyIntMap g == h@ then the following holds:
--
-- @
-- map ('fromVertexKL' h) ('Data.Graph.vertices' $ 'toGraphKL' h)                               == 'Data.IntSet.toAscList' ('Algebra.Graph.AdjacencyIntMap.vertexIntSet' g)
-- map (\\(x, y) -> ('fromVertexKL' h x, 'fromVertexKL' h y)) ('Data.Graph.edges' $ 'toGraphKL' h) == 'Algebra.Graph.AdjacencyIntMap.edgeList' g
-- 'toGraphKL' (fromAdjacencyIntMap (1 * 2 + 3 * 1))                             == 'array' (0,2) [(0,[1]), (1,[]), (2,[0])]
-- 'toGraphKL' (fromAdjacencyIntMap (1 * 2 + 2 * 1))                             == 'array' (0,1) [(0,[1]), (1,[0])]
-- @
fromAdjacencyIntMap :: AIM.AdjacencyIntMap -> GraphKL Int
fromAdjacencyIntMap :: AdjacencyIntMap -> GraphKL Int
fromAdjacencyIntMap AdjacencyIntMap
aim = GraphKL
    { toGraphKL :: Graph
toGraphKL    = Graph
g
    , fromVertexKL :: Int -> Int
fromVertexKL = \Int
x -> case Int -> ((), Int, [Int])
r Int
x of (()
_, Int
v, [Int]
_) -> Int
v
    , toVertexKL :: Int -> Maybe Int
toVertexKL   = Int -> Maybe Int
t }
  where
    (Graph
g, Int -> ((), Int, [Int])
r, Int -> Maybe Int
t) = [((), Int, [Int])]
-> (Graph, Int -> ((), Int, [Int]), Int -> Maybe Int)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Int -> (node, key, [key]), key -> Maybe Int)
KL.graphFromEdges [ ((), Int
x, [Int]
ys) | (Int
x, [Int]
ys) <- AdjacencyIntMap -> [(Int, [Int])]
AIM.adjacencyList AdjacencyIntMap
aim ]

-- | Compute the /depth-first search/ forest of a graph.
--
-- In the following examples we will use the helper function:
--
-- @
-- (%) :: Ord a => ('GraphKL' a -> b) -> 'AM.AdjacencyMap' a -> b
-- f % x = f ('fromAdjacencyMap' x)
-- @
--
-- for greater clarity.
--
-- @
-- 'AM.forest' (dfsForest % 'AM.edge' 1 1)           == 'AM.vertex' 1
-- 'AM.forest' (dfsForest % 'AM.edge' 1 2)           == 'AM.edge' 1 2
-- 'AM.forest' (dfsForest % 'AM.edge' 2 1)           == 'AM.vertices' [1,2]
-- 'AM.isSubgraphOf' ('AM.forest' $ dfsForest % x) x == True
-- dfsForest % 'AM.forest' (dfsForest % x)      == dfsForest % x
-- dfsForest % 'AM.vertices' vs                 == 'map' (\\v -> Node v []) ('Data.List.nub' $ 'Data.List.sort' vs)
-- dfsForest % (3 * (1 + 4) * (1 + 5))     == [ Node { rootLabel = 1
--                                                   , subForest = [ Node { rootLabel = 5
--                                                                        , subForest = [] }]}
--                                            , Node { rootLabel = 3
--                                                   , subForest = [ Node { rootLabel = 4
--                                                                        , subForest = [] }]}]
-- @
dfsForest :: GraphKL a -> Forest a
dfsForest :: forall a. GraphKL a -> Forest a
dfsForest (GraphKL Graph
g Int -> a
r a -> Maybe Int
_) = (Tree Int -> Tree a) -> [Tree Int] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> a) -> Tree Int -> Tree a
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> a
r) (Graph -> [Tree Int]
KL.dff Graph
g)

-- | Compute the /depth-first search/ forest of a graph, searching from each of
-- the given vertices in order. Note that the resulting forest does not
-- necessarily span the whole graph, as some vertices may be unreachable.
--
-- In the following examples we will use the helper function:
--
-- @
-- (%) :: Ord a => ('GraphKL' a -> b) -> 'AM.AdjacencyMap' a -> b
-- f % x = f ('fromAdjacencyMap' x)
-- @
--
-- for greater clarity.
--
-- @
-- 'AM.forest' $ (dfsForestFrom % 'AM.edge' 1 1) [1]          == 'AM.vertex' 1
-- 'AM.forest' $ (dfsForestFrom % 'AM.edge' 1 2) [0]          == 'AM.empty'
-- 'AM.forest' $ (dfsForestFrom % 'AM.edge' 1 2) [1]          == 'AM.edge' 1 2
-- 'AM.forest' $ (dfsForestFrom % 'AM.edge' 1 2) [2]          == 'AM.vertex' 2
-- 'AM.forest' $ (dfsForestFrom % 'AM.edge' 1 2) [2,1]        == 'AM.vertices' [1,2]
-- 'AM.isSubgraphOf' ('AM.forest' $ dfsForestFrom % x $ vs) x == True
-- dfsForestFrom % x $ 'AM.vertexList' x                 == 'dfsForest' % x
-- dfsForestFrom % 'AM.vertices' vs $ vs                 == 'map' (\\v -> Node v []) ('Data.List.nub' vs)
-- dfsForestFrom % x $ []                           == []
-- dfsForestFrom % (3 * (1 + 4) * (1 + 5)) $ [1,4]  == [ Node { rootLabel = 1
--                                                            , subForest = [ Node { rootLabel = 5
--                                                                                 , subForest = [] }
--                                                     , Node { rootLabel = 4
--                                                            , subForest = [] }]
-- @
dfsForestFrom :: GraphKL a -> [a] -> Forest a
dfsForestFrom :: forall a. GraphKL a -> [a] -> Forest a
dfsForestFrom (GraphKL Graph
g Int -> a
r a -> Maybe Int
t) = (Tree Int -> Tree a) -> [Tree Int] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> a) -> Tree Int -> Tree a
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> a
r) ([Tree Int] -> [Tree a]) -> ([a] -> [Tree Int]) -> [a] -> [Tree a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph -> [Int] -> [Tree Int]
KL.dfs Graph
g ([Int] -> [Tree Int]) -> ([a] -> [Int]) -> [a] -> [Tree Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe Int) -> [a] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe a -> Maybe Int
t

-- | Compute the list of vertices visited by the /depth-first search/ in a
-- graph, when searching from each of the given vertices in order.
--
-- In the following examples we will use the helper function:
--
-- @
-- (%) :: Ord a => ('GraphKL' a -> b) -> 'AM.AdjacencyMap' a -> b
-- f % x = f ('fromAdjacencyMap' x)
-- @
--
-- for greater clarity.
--
-- @
-- dfs % 'AM.edge' 1 1 $ [1]   == [1]
-- dfs % 'AM.edge' 1 2 $ [0]   == []
-- dfs % 'AM.edge' 1 2 $ [1]   == [1,2]
-- dfs % 'AM.edge' 1 2 $ [2]   == [2]
-- dfs % 'AM.edge' 1 2 $ [1,2] == [1,2]
-- dfs % 'AM.edge' 1 2 $ [2,1] == [2,1]
-- dfs % x        $ []    == []
--
-- dfs % (3 * (1 + 4) * (1 + 5)) $ [1,4]     == [1,5,4]
-- 'Data.List.and' [ 'AM.hasVertex' v x | v <- dfs % x $ vs ] == True
-- @
dfs :: GraphKL a -> [a] -> [a]
dfs :: forall a. GraphKL a -> [a] -> [a]
dfs GraphKL a
x = (Tree a -> [a]) -> [Tree a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [a]
forall a. Tree a -> [a]
flatten ([Tree a] -> [a]) -> ([a] -> [Tree a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GraphKL a -> [a] -> [Tree a]
forall a. GraphKL a -> [a] -> Forest a
dfsForestFrom GraphKL a
x

-- | Compute the /topological sort/ of a graph. Note that this function returns
-- a result even if the graph is cyclic.
--
-- In the following examples we will use the helper function:
--
-- @
-- (%) :: Ord a => ('GraphKL' a -> b) -> 'AM.AdjacencyMap' a -> b
-- f % x = f ('fromAdjacencyMap' x)
-- @
--
-- for greater clarity.
--
-- @
-- topSort % (1 * 2 + 3 * 1) == [3,1,2]
-- topSort % (1 * 2 + 2 * 1) == [1,2]
-- @
topSort :: GraphKL a -> [a]
topSort :: forall a. GraphKL a -> [a]
topSort (GraphKL Graph
g Int -> a
r a -> Maybe Int
_) = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Int -> a
r (Graph -> [Int]
KL.topSort Graph
g)

-- TODO: Add docs and tests.
scc :: Ord a => AM.AdjacencyMap a -> AM.AdjacencyMap (NonEmpty.AdjacencyMap a)
scc :: forall a. Ord a => AdjacencyMap a -> AdjacencyMap (AdjacencyMap a)
scc AdjacencyMap a
m = (Int -> AdjacencyMap a)
-> AdjacencyMap Int -> AdjacencyMap (AdjacencyMap a)
forall a b.
(Ord a, Ord b) =>
(a -> b) -> AdjacencyMap a -> AdjacencyMap b
AM.gmap (Map Int (AdjacencyMap a)
component Map Int (AdjacencyMap a) -> Int -> AdjacencyMap a
forall k a. Ord k => Map k a -> k -> a
Map.!) (AdjacencyMap Int -> AdjacencyMap (AdjacencyMap a))
-> AdjacencyMap Int -> AdjacencyMap (AdjacencyMap a)
forall a b. (a -> b) -> a -> b
$ AdjacencyMap Int -> AdjacencyMap Int
forall a. Ord a => AdjacencyMap a -> AdjacencyMap a
removeSelfLoops (AdjacencyMap Int -> AdjacencyMap Int)
-> AdjacencyMap Int -> AdjacencyMap Int
forall a b. (a -> b) -> a -> b
$ (a -> Int) -> AdjacencyMap a -> AdjacencyMap Int
forall a b.
(Ord a, Ord b) =>
(a -> b) -> AdjacencyMap a -> AdjacencyMap b
AM.gmap (Map a Int
leader Map a Int -> a -> Int
forall k a. Ord k => Map k a -> k -> a
Map.!) AdjacencyMap a
m
  where
    GraphKL Graph
g Int -> a
decode a -> Maybe Int
_ = AdjacencyMap a -> GraphKL a
forall a. Ord a => AdjacencyMap a -> GraphKL a
fromAdjacencyMap AdjacencyMap a
m
    sccs :: [[Int]]
sccs      = (Tree Int -> [Int]) -> [Tree Int] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map Tree Int -> [Int]
forall a. Tree a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Graph -> [Tree Int]
KL.scc Graph
g)
    leader :: Map a Int
leader    = [(a, Int)] -> Map a Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (Int -> a
decode Int
y, Int
x)      | Int
x:[Int]
xs <- [[Int]]
sccs, Int
y <- Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
xs ]
    component :: Map Int (AdjacencyMap a)
component = [(Int, AdjacencyMap a)] -> Map Int (AdjacencyMap a)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (Int
x, [Int] -> AdjacencyMap a
expand (Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
xs)) | Int
x:[Int]
xs <- [[Int]]
sccs ]
    expand :: [Int] -> AdjacencyMap a
expand [Int]
xs = Maybe (AdjacencyMap a) -> AdjacencyMap a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (AdjacencyMap a) -> AdjacencyMap a)
-> Maybe (AdjacencyMap a) -> AdjacencyMap a
forall a b. (a -> b) -> a -> b
$ AdjacencyMap a -> Maybe (AdjacencyMap a)
forall a. AdjacencyMap a -> Maybe (AdjacencyMap a)
NonEmpty.toNonEmpty (AdjacencyMap a -> Maybe (AdjacencyMap a))
-> AdjacencyMap a -> Maybe (AdjacencyMap a)
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> AdjacencyMap a -> AdjacencyMap a
forall a. (a -> Bool) -> AdjacencyMap a -> AdjacencyMap a
AM.induce (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a
s) AdjacencyMap a
m
      where
        s :: Set a
s = [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ((Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Int -> a
decode [Int]
xs)

removeSelfLoops :: Ord a => AM.AdjacencyMap a -> AM.AdjacencyMap a
removeSelfLoops :: forall a. Ord a => AdjacencyMap a -> AdjacencyMap a
removeSelfLoops AdjacencyMap a
m = (a -> AdjacencyMap a -> AdjacencyMap a)
-> AdjacencyMap a -> [a] -> AdjacencyMap 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
x -> a -> a -> AdjacencyMap a -> AdjacencyMap a
forall a. Ord a => a -> a -> AdjacencyMap a -> AdjacencyMap a
AM.removeEdge a
x a
x) AdjacencyMap a
m (AdjacencyMap a -> [a]
forall a. AdjacencyMap a -> [a]
AM.vertexList AdjacencyMap a
m)