-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | Some functions missing from "Data.Graph".
module Util.Graphs where
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Array.IArray as IArray
import Data.Array.IArray ( (!), (//) )
import Data.Graph
import qualified Data.Tree as Tree

import qualified Util.Array as Array
import qualified Util.Lists as Lists


build :: [Edge] -> Graph
build :: [Edge] -> Graph
build [Edge]
edges = Edge -> [Edge] -> Graph
buildG (Vertex
0, Vertex
upper) [Edge]
unique
    where
    unique :: [Edge]
unique = forall a. Ord a => [a] -> [a]
Lists.unique [Edge]
edges
    upper :: Vertex
upper = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (-Vertex
1 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (\(Vertex
x, Vertex
y) -> forall a. Ord a => a -> a -> a
max Vertex
x Vertex
y) [Edge]
unique)

-- | Roots are all vertices with no parents.
roots_of :: Graph -> [Vertex]
roots_of :: Graph -> [Vertex]
roots_of Graph
graph =
    forall a. Set a -> [a]
Set.toList forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Ord a => a -> Set a -> Set a
Set.delete) (forall a. Ord a => [a] -> Set a
Set.fromList [Vertex
lo..Vertex
hi])
        (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
IArray.elems Graph
graph))
    where (Vertex
lo, Vertex
hi) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Graph
graph

-- | This is like 'dfs', except it allows duplicated vertices.  So don't
-- use it one a graph with cycles.
to_forest :: Graph -> Tree.Forest Vertex
to_forest :: Graph -> Forest Vertex
to_forest Graph
graph = forall a b. (a -> b) -> [a] -> [b]
map (forall {a :: * -> * -> *} {a}.
(IArray a [a], Ix a) =>
a a [a] -> a -> Tree a
generate Graph
graph) (Graph -> [Vertex]
roots_of Graph
graph)
    where
    generate :: a a [a] -> a -> Tree a
generate a a [a]
graph a
vertex = forall a. a -> [Tree a] -> Tree a
Node a
vertex (forall a b. (a -> b) -> [a] -> [b]
map (a a [a] -> a -> Tree a
generate a a [a]
graph) (a a [a]
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!a
vertex))

draw :: Graph -> String
draw :: Graph -> String
draw = [Tree String] -> String
Tree.drawForest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Show a => a -> String
show) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph -> Forest Vertex
to_forest

-- | Remove the edge if it already exists, create a new one of it doesn't.
-- Return Nothing if adding an edge would create a cycle.
toggle_edge :: Edge -> Graph -> Maybe Graph
toggle_edge :: Edge -> Graph -> Maybe Graph
toggle_edge Edge
edge Graph
graph
    | Graph -> Edge -> Bool
has_edge Graph
graph Edge
edge = forall a. a -> Maybe a
Just ([Edge] -> Graph -> Graph
remove_edges [Edge
edge] Graph
graph)
    | Edge -> Graph -> Bool
would_make_cycle Edge
edge Graph
graph = forall a. Maybe a
Nothing
    | Bool
otherwise = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [Edge] -> Graph -> Graph
add_edges [Edge
edge] Graph
graph

-- | Splice @new@ into the graph above @to@.  The parents of @to@ are detached
-- from it and re-attached to @new@.  Then @new@ is attached above @to@.
--
-- This operation should be idempotent.
splice_above :: Vertex -> Vertex -> Graph -> Graph
splice_above :: Vertex -> Vertex -> Graph -> Graph
splice_above Vertex
new Vertex
to Graph
graph =
    -- If I don't filter p/=new, a duplicate splice will cause a vertex to
    -- loop back to itself.
    [Edge] -> Graph -> Graph
add_edges ((Vertex
new, Vertex
to) forall a. a -> [a] -> [a]
: [(Vertex
p, Vertex
new) | Vertex
p <- [Vertex]
ps, Vertex
p forall a. Eq a => a -> a -> Bool
/= Vertex
new]) forall a b. (a -> b) -> a -> b
$
        [Edge] -> Graph -> Graph
remove_edges [(Vertex
p, Vertex
to) | Vertex
p <- [Vertex]
ps] Graph
graph
    where ps :: [Vertex]
ps = Graph -> Vertex -> [Vertex]
parents Graph
graph Vertex
to

-- | Splice @new@ into the graph below @to@.  The children of @to@ are
-- detached and re-attached to @new@.  Then @to@ is attached above @new@.
--
-- This operation should be idempotent.
splice_below :: Vertex -> Vertex -> Graph -> Graph
splice_below :: Vertex -> Vertex -> Graph -> Graph
splice_below Vertex
new Vertex
to Graph
graph =
    [Edge] -> Graph -> Graph
add_edges ((Vertex
to, Vertex
new) forall a. a -> [a] -> [a]
: [(Vertex
new, Vertex
c) | Vertex
c <- [Vertex]
cs, Vertex
c forall a. Eq a => a -> a -> Bool
/= Vertex
new]) forall a b. (a -> b) -> a -> b
$
        [Edge] -> Graph -> Graph
remove_edges (forall a b. (a -> b) -> [a] -> [b]
map ((,) Vertex
to) [Vertex]
cs) Graph
graph
    where cs :: [Vertex]
cs = Graph -> Vertex -> [Vertex]
children Graph
graph Vertex
to

-- | Get the parents of a Vertex.
parents :: Graph -> Vertex -> [Vertex]
parents :: Graph -> Vertex -> [Vertex]
parents Graph
graph Vertex
v = [Vertex
p | (Vertex
p, [Vertex]
cs) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
IArray.assocs Graph
graph, Vertex
v forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Vertex]
cs]

-- | Get the children of a Vertex.
children :: Graph -> Vertex -> [Vertex]
children :: Graph -> Vertex -> [Vertex]
children Graph
graph Vertex
v
    | forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
v Graph
graph = Graph
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
v
    | Bool
otherwise = []

would_make_cycle :: Edge -> Graph -> Bool
would_make_cycle :: Edge -> Graph -> Bool
would_make_cycle (Vertex
from, Vertex
to) Graph
graph =
    Vertex
from forall a. Eq a => a -> a -> Bool
== Vertex
to
    Bool -> Bool -> Bool
|| forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
from Graph
graph Bool -> Bool -> Bool
&& forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
to Graph
graph
        Bool -> Bool -> Bool
&& Graph -> Vertex -> Vertex -> Bool
path Graph
graph Vertex
to Vertex
from

has_cycle :: Graph -> Bool
has_cycle :: Graph -> Bool
has_cycle Graph
graph = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> [Tree a]
Tree.subForest) (Graph -> Forest Vertex
scc Graph
graph)
    Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Eq a => a -> a -> Bool
(==)) (Graph -> [Edge]
edges Graph
graph)

has_edge :: Graph -> Edge -> Bool
has_edge :: Graph -> Edge -> Bool
has_edge Graph
graph (Vertex
from, Vertex
to) = forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
from Graph
graph Bool -> Bool -> Bool
&& Vertex
to forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Graph
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
from

-- | A lonely vertex has no edges.
lonely_vertex :: Graph -> Vertex -> Bool
lonely_vertex :: Graph -> Vertex -> Bool
lonely_vertex Graph
graph Vertex
vertex =
    Bool -> Bool
not (forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
vertex Graph
graph) Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Graph
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
vertex)

add_edges :: [Edge] -> Graph -> Graph
add_edges :: [Edge] -> Graph -> Graph
add_edges [Edge]
edges Graph
graph =
    forall (a :: * -> * -> *) e i e'.
(IArray a e, Ix i) =>
(e -> e' -> e) -> e -> (i, i) -> [(i, e')] -> a i e
IArray.accumArray forall {a}. Eq a => [a] -> a -> [a]
add [] Edge
new_bounds
        ([Edge]
edges forall a. [a] -> [a] -> [a]
++ [(Vertex
p, Vertex
c) | (Vertex
p, [Vertex]
cs) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
IArray.assocs Graph
graph, Vertex
c <- [Vertex]
cs])
    where
    (Vertex
low, Vertex
high) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Graph
graph
    flattened :: [Vertex]
flattened = [Vertex
v | (Vertex
p, Vertex
c) <- [Edge]
edges, Vertex
v <- [Vertex
p, Vertex
c]]
    new_bounds :: Edge
new_bounds = (forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (Vertex
low forall a. a -> [a] -> [a]
: [Vertex]
flattened), forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Vertex
high forall a. a -> [a] -> [a]
: [Vertex]
flattened))
    add :: [a] -> a -> [a]
add [a]
cs a
c = if a
c forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a]
cs then [a]
cs else a
cforall a. a -> [a] -> [a]
:[a]
cs

remove_edges :: [Edge] -> Graph -> Graph
remove_edges :: [Edge] -> Graph -> Graph
remove_edges [Edge]
edges Graph
graph
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Edge]
edges = Graph
graph
    | Bool
otherwise =
        Graph
graph forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
// [(Vertex
from, forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [Edge]
groups) (Graph
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
from))
            | (Vertex
from, [Edge]
groups) <- [(Vertex, [Edge])]
grouped]
    where
    in_bounds :: [Edge]
in_bounds = forall a. (a -> Bool) -> [a] -> [a]
filter ((\Vertex
p -> forall a. Vertex -> Array a -> Bool
Array.in_bounds Vertex
p Graph
graph) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [Edge]
edges
    grouped :: [(Vertex, [Edge])]
grouped = forall key a. Ord key => (a -> key) -> [a] -> [(key, [a])]
Lists.keyedGroupSort forall a b. (a, b) -> a
fst [Edge]
in_bounds

-- | Increment all vertices at and above, insert new empty vertex.
insert_vertex :: Int -> Graph -> Graph
insert_vertex :: Vertex -> Graph -> Graph
insert_vertex Vertex
vertex Graph
graph = (Vertex -> Vertex) -> Graph -> Graph
map_vertices Vertex -> Vertex
incr Graph
graph
    where incr :: Vertex -> Vertex
incr Vertex
v = if Vertex
v forall a. Ord a => a -> a -> Bool
< Vertex
vertex then Vertex
v else Vertex
vforall a. Num a => a -> a -> a
+Vertex
1

-- | Remove a vertex.  All vertices pointing to the removed vertex instead
-- point to what pointed to it.
remove_vertex :: Int -> Graph -> Graph
remove_vertex :: Vertex -> Graph -> Graph
remove_vertex Vertex
vertex Graph
graph
    | Vertex
vertex forall a. Vertex -> Array a -> Bool
`Array.in_bounds` Graph
graph =
        (Vertex -> Vertex) -> Graph -> Graph
map_vertices Vertex -> Vertex
decr (Vertex -> Graph -> Graph
unlink_vertex Vertex
vertex Graph
graph)
    | Bool
otherwise = Graph
graph
    where decr :: Vertex -> Vertex
decr Vertex
v = if Vertex
v forall a. Ord a => a -> a -> Bool
> Vertex
vertex then Vertex
vforall a. Num a => a -> a -> a
-Vertex
1 else Vertex
v

-- | All vertices pointing to the removed vertex instead point to what pointed
-- to it.  It will be removed from the list of roots.
unlink_vertex :: Int -> Graph -> Graph
unlink_vertex :: Vertex -> Graph -> Graph
unlink_vertex Vertex
vertex Graph
graph =
    forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
IArray.amap (forall a. Eq a => a -> [a] -> [a] -> [a]
Lists.replace1 Vertex
v (Graph
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
v)) Graph
graph forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
// [(Vertex
v, [])]
    where v :: Vertex
v = forall a. String -> Vertex -> Array a -> Vertex
Array.assert_in_bounds String
"unlink_vertex" Vertex
vertex Graph
graph


-- | Transform all the vertices by the given function.  If multiple vertices
-- are transformed to the same value, the one with the originally highest
-- vertex wins.
map_vertices :: (Vertex -> Vertex) -> Graph -> Graph
map_vertices :: (Vertex -> Vertex) -> Graph -> Graph
map_vertices Vertex -> Vertex
f Graph
graph = forall a. [a] -> Array a
Array.from_list forall a b. (a -> b) -> a -> b
$
    forall a. a -> [(Vertex, a)] -> [a]
strip_indices [] forall a b. (a -> b) -> a -> b
$ forall k a. Eq k => (a -> k) -> [a] -> [a]
Lists.dropInitialDups forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$
        forall a b. (a -> b) -> [a] -> [b]
map (\(Vertex
p, [Vertex]
cs) -> (Vertex -> Vertex
f Vertex
p, forall a b. (a -> b) -> [a] -> [b]
map Vertex -> Vertex
f [Vertex]
cs)) (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
IArray.assocs Graph
graph)

strip_indices :: a -> [(Int, a)] -> [a]
strip_indices :: forall a. a -> [(Vertex, a)] -> [a]
strip_indices a
def = forall {t}. (Ord t, Num t) => t -> [(t, a)] -> [a]
go Vertex
0
    where
    go :: t -> [(t, a)] -> [a]
go t
_ [] = []
    go t
prev lst :: [(t, a)]
lst@((t
i, a
v):[(t, a)]
xs)
        | t
prev forall a. Ord a => a -> a -> Bool
< t
i = a
def forall a. a -> [a] -> [a]
: t -> [(t, a)] -> [a]
go (t
prevforall a. Num a => a -> a -> a
+t
1) [(t, a)]
lst
        | Bool
otherwise = a
v forall a. a -> [a] -> [a]
: t -> [(t, a)] -> [a]
go (t
prevforall a. Num a => a -> a -> a
+t
1) [(t, a)]
xs

-- | Move a vertex.  The graph remains the same, but the @from@ vertex number
-- will be changed to @to@ and vice versa.
move :: Vertex -> Vertex -> Graph -> Maybe Graph
move :: Vertex -> Vertex -> Graph -> Maybe Graph
move Vertex
from Vertex
to Graph
graph = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {a :: * -> * -> *}.
IArray a [Vertex] =>
a Vertex [Vertex] -> a Vertex [Vertex]
move forall a b. (a -> b) -> a -> b
$ forall a. Vertex -> Array a -> Maybe (Array a)
Array.check Vertex
to forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. Vertex -> Array a -> Maybe (Array a)
Array.check Vertex
from Graph
graph
    where
    -- Swap array elements.  Then swap all referents
    move :: a Vertex [Vertex] -> a Vertex [Vertex]
move a Vertex [Vertex]
graph
        | Vertex
from forall a. Eq a => a -> a -> Bool
== Vertex
to = a Vertex [Vertex]
graph
        | Bool
otherwise = forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
IArray.amap (forall a b. (a -> b) -> [a] -> [b]
map Vertex -> Vertex
relink) forall a b. (a -> b) -> a -> b
$
            a Vertex [Vertex]
graph forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
// [(Vertex
from, a Vertex [Vertex]
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
to), (Vertex
to, a Vertex [Vertex]
graphforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
from)]
    relink :: Vertex -> Vertex
relink Vertex
v
        | Vertex
v forall a. Eq a => a -> a -> Bool
== Vertex
from = Vertex
to
        | Vertex
v forall a. Eq a => a -> a -> Bool
== Vertex
to = Vertex
from
        | Bool
otherwise = Vertex
v