-- | Wrapper around Data.Graph with support for edge labels
{-# LANGUAGE ScopedTypeVariables #-}
module Distribution.Solver.Modular.LabeledGraph (
    -- * Graphs
    Graph
  , Vertex
    -- ** Building graphs
  , graphFromEdges
  , graphFromEdges'
  , buildG
  , transposeG
    -- ** Graph properties
  , vertices
  , edges
    -- ** Operations on the underlying unlabeled graph
  , forgetLabels
  , topSort
  ) where

import Distribution.Solver.Compat.Prelude
import Prelude ()

import Data.Array
import Data.Graph (Vertex, Bounds)
import qualified Data.Graph as G

{-------------------------------------------------------------------------------
  Types
-------------------------------------------------------------------------------}

type Graph e = Array Vertex [(e, Vertex)]
type Edge  e = (Vertex, e, Vertex)

{-------------------------------------------------------------------------------
  Building graphs
-------------------------------------------------------------------------------}

-- | Construct an edge-labeled graph
--
-- This is a simple adaptation of the definition in Data.Graph
graphFromEdges :: forall key node edge. Ord key
               => [ (node, key, [(edge, key)]) ]
               -> ( Graph edge
                  , Vertex -> (node, key, [(edge, key)])
                  , key -> Maybe Vertex
                  )
graphFromEdges :: forall key node edge.
Ord key =>
[(node, key, [(edge, key)])]
-> (Graph edge, Vertex -> (node, key, [(edge, key)]),
    key -> Maybe Vertex)
graphFromEdges [(node, key, [(edge, key)])]
edges0 =
    (Array Vertex [(edge, Vertex)]
graph, \Vertex
v -> Array Vertex (node, key, [(edge, key)])
vertex_map forall i e. Ix i => Array i e -> i -> e
! Vertex
v, key -> Maybe Vertex
key_vertex)
  where
    max_v :: Vertex
max_v        = forall (t :: * -> *) a. Foldable t => t a -> Vertex
length [(node, key, [(edge, key)])]
edges0 forall a. Num a => a -> a -> a
- Vertex
1
    bounds0 :: (Vertex, Vertex)
bounds0      = (Vertex
0, Vertex
max_v) :: (Vertex, Vertex)
    sorted_edges :: [(node, key, [(edge, key)])]
sorted_edges = forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy forall {a} {a} {c} {a} {c}.
Ord a =>
(a, a, c) -> (a, a, c) -> Ordering
lt [(node, key, [(edge, key)])]
edges0
    edges1 :: [(Vertex, (node, key, [(edge, key)]))]
edges1       = forall a b. [a] -> [b] -> [(a, b)]
zip [Vertex
0..] [(node, key, [(edge, key)])]
sorted_edges

    graph :: Array Vertex [(edge, Vertex)]
graph        = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Vertex, Vertex)
bounds0 [(Vertex
v, (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (edge, key) -> Maybe (edge, Vertex)
mk_edge [(edge, key)]
ks))
                                 | (Vertex
v, (node
_, key
_, [(edge, key)]
ks)) <- [(Vertex, (node, key, [(edge, key)]))]
edges1]
    key_map :: Array Vertex key
key_map      = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Vertex, Vertex)
bounds0 [(Vertex
v, key
k                    )
                                 | (Vertex
v, (node
_, key
k, [(edge, key)]
_ )) <- [(Vertex, (node, key, [(edge, key)]))]
edges1]
    vertex_map :: Array Vertex (node, key, [(edge, key)])
vertex_map   = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Vertex, Vertex)
bounds0 [(Vertex, (node, key, [(edge, key)]))]
edges1

    (a
_,a
k1,c
_) lt :: (a, a, c) -> (a, a, c) -> Ordering
`lt` (a
_,a
k2,c
_) = a
k1 forall a. Ord a => a -> a -> Ordering
`compare` a
k2

    mk_edge :: (edge, key) -> Maybe (edge, Vertex)
    mk_edge :: (edge, key) -> Maybe (edge, Vertex)
mk_edge (edge
edge, key
key) = do Vertex
v <- key -> Maybe Vertex
key_vertex key
key ; forall (m :: * -> *) a. Monad m => a -> m a
return (edge
edge, Vertex
v)

    --  returns Nothing for non-interesting vertices
    key_vertex :: key -> Maybe Vertex
    key_vertex :: key -> Maybe Vertex
key_vertex key
k = Vertex -> Vertex -> Maybe Vertex
findVertex Vertex
0 Vertex
max_v
      where
        findVertex :: Vertex -> Vertex -> Maybe Vertex
findVertex Vertex
a Vertex
b
          | Vertex
a forall a. Ord a => a -> a -> Bool
> Vertex
b     = forall a. Maybe a
Nothing
          | Bool
otherwise = case forall a. Ord a => a -> a -> Ordering
compare key
k (Array Vertex key
key_map forall i e. Ix i => Array i e -> i -> e
! Vertex
mid) of
              Ordering
LT -> Vertex -> Vertex -> Maybe Vertex
findVertex Vertex
a (Vertex
midforall a. Num a => a -> a -> a
-Vertex
1)
              Ordering
EQ -> forall a. a -> Maybe a
Just Vertex
mid
              Ordering
GT -> Vertex -> Vertex -> Maybe Vertex
findVertex (Vertex
midforall a. Num a => a -> a -> a
+Vertex
1) Vertex
b
          where
            mid :: Vertex
mid = Vertex
a forall a. Num a => a -> a -> a
+ (Vertex
b forall a. Num a => a -> a -> a
- Vertex
a) forall a. Integral a => a -> a -> a
`div` Vertex
2

graphFromEdges' :: Ord key
                => [ (node, key, [(edge, key)]) ]
                -> ( Graph edge
                   , Vertex -> (node, key, [(edge, key)])
                   )
graphFromEdges' :: forall key node edge.
Ord key =>
[(node, key, [(edge, key)])]
-> (Graph edge, Vertex -> (node, key, [(edge, key)]))
graphFromEdges' [(node, key, [(edge, key)])]
x = (Graph edge
a,Vertex -> (node, key, [(edge, key)])
b)
  where
    (Graph edge
a,Vertex -> (node, key, [(edge, key)])
b,key -> Maybe Vertex
_) = forall key node edge.
Ord key =>
[(node, key, [(edge, key)])]
-> (Graph edge, Vertex -> (node, key, [(edge, key)]),
    key -> Maybe Vertex)
graphFromEdges [(node, key, [(edge, key)])]
x

transposeG :: Graph e -> Graph e
transposeG :: forall e. Graph e -> Graph e
transposeG Graph e
g = forall e. (Vertex, Vertex) -> [Edge e] -> Graph e
buildG (forall i e. Array i e -> (i, i)
bounds Graph e
g) (forall e. Graph e -> [Edge e]
reverseE Graph e
g)

buildG :: Bounds -> [Edge e] -> Graph e
buildG :: forall e. (Vertex, Vertex) -> [Edge e] -> Graph e
buildG (Vertex, Vertex)
bounds0 [Edge e]
edges0 = forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray (forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] (Vertex, Vertex)
bounds0 (forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a} {b}. (a, a, b) -> (a, (a, b))
reassoc [Edge e]
edges0)
  where
    reassoc :: (a, a, b) -> (a, (a, b))
reassoc (a
v, a
e, b
w) = (a
v, (a
e, b
w))

reverseE :: Graph e -> [Edge e]
reverseE :: forall e. Graph e -> [Edge e]
reverseE Graph e
g = [ (Vertex
w, e
e, Vertex
v) | (Vertex
v, e
e, Vertex
w) <- forall e. Graph e -> [Edge e]
edges Graph e
g ]

{-------------------------------------------------------------------------------
  Graph properties
-------------------------------------------------------------------------------}

vertices :: Graph e -> [Vertex]
vertices :: forall e. Graph e -> [Vertex]
vertices = forall i e. Ix i => Array i e -> [i]
indices

edges :: Graph e -> [Edge e]
edges :: forall e. Graph e -> [Edge e]
edges Graph e
g = [ (Vertex
v, e
e, Vertex
w) | Vertex
v <- forall e. Graph e -> [Vertex]
vertices Graph e
g, (e
e, Vertex
w) <- Graph e
gforall i e. Ix i => Array i e -> i -> e
!Vertex
v ]

{-------------------------------------------------------------------------------
  Operations on the underlying unlabelled graph
-------------------------------------------------------------------------------}

forgetLabels :: Graph e -> G.Graph
forgetLabels :: forall e. Graph e -> Graph
forgetLabels = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd)

topSort :: Graph e -> [Vertex]
topSort :: forall e. Graph e -> [Vertex]
topSort = Graph -> [Vertex]
G.topSort forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e. Graph e -> Graph
forgetLabels