-- | 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 Array Vertex (node, key, [(edge, key)])
-> Vertex -> (node, key, [(edge, key)])
forall i e. Ix i => Array i e -> i -> e
! Vertex
v, key -> Maybe Vertex
key_vertex)
  where
    max_v :: Vertex
max_v        = [(node, key, [(edge, key)])] -> Vertex
forall a. [a] -> Vertex
forall (t :: * -> *) a. Foldable t => t a -> Vertex
length [(node, key, [(edge, key)])]
edges0 Vertex -> Vertex -> Vertex
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 = ((node, key, [(edge, key)])
 -> (node, key, [(edge, key)]) -> Ordering)
-> [(node, key, [(edge, key)])] -> [(node, key, [(edge, key)])]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (node, key, [(edge, key)])
-> (node, key, [(edge, key)]) -> Ordering
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       = [Vertex]
-> [(node, key, [(edge, key)])]
-> [(Vertex, (node, key, [(edge, key)]))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vertex
0..] [(node, key, [(edge, key)])]
sorted_edges

    graph :: Array Vertex [(edge, Vertex)]
graph        = (Vertex, Vertex)
-> [(Vertex, [(edge, Vertex)])] -> Array Vertex [(edge, Vertex)]
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Vertex, Vertex)
bounds0 [(Vertex
v, (((edge, key) -> Maybe (edge, Vertex))
-> [(edge, key)] -> [(edge, Vertex)]
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      = (Vertex, Vertex) -> [(Vertex, key)] -> Array Vertex key
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   = (Vertex, Vertex)
-> [(Vertex, (node, key, [(edge, key)]))]
-> Array Vertex (node, key, [(edge, key)])
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 a -> a -> Ordering
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 ; (edge, Vertex) -> Maybe (edge, Vertex)
forall a. a -> Maybe a
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 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
b     = Maybe Vertex
forall a. Maybe a
Nothing
          | Bool
otherwise = case key -> key -> Ordering
forall a. Ord a => a -> a -> Ordering
compare key
k (Array Vertex key
key_map Array Vertex key -> Vertex -> key
forall i e. Ix i => Array i e -> i -> e
! Vertex
mid) of
              Ordering
LT -> Vertex -> Vertex -> Maybe Vertex
findVertex Vertex
a (Vertex
midVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
-Vertex
1)
              Ordering
EQ -> Vertex -> Maybe Vertex
forall a. a -> Maybe a
Just Vertex
mid
              Ordering
GT -> Vertex -> Vertex -> Maybe Vertex
findVertex (Vertex
midVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
1) Vertex
b
          where
            mid :: Vertex
mid = Vertex
a Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ (Vertex
b Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
a) Vertex -> Vertex -> Vertex
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
_) = [(node, key, [(edge, key)])]
-> (Graph edge, Vertex -> (node, key, [(edge, key)]),
    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 = (Vertex, Vertex) -> [Edge e] -> Graph e
forall e. (Vertex, Vertex) -> [Edge e] -> Graph e
buildG (Graph e -> (Vertex, Vertex)
forall i e. Array i e -> (i, i)
bounds Graph e
g) (Graph e -> [Edge e]
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 = ([(e, Vertex)] -> (e, Vertex) -> [(e, Vertex)])
-> [(e, Vertex)]
-> (Vertex, Vertex)
-> [(Vertex, (e, Vertex))]
-> Array Vertex [(e, Vertex)]
forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray (((e, Vertex) -> [(e, Vertex)] -> [(e, Vertex)])
-> [(e, Vertex)] -> (e, Vertex) -> [(e, Vertex)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] (Vertex, Vertex)
bounds0 ((Edge e -> (Vertex, (e, Vertex)))
-> [Edge e] -> [(Vertex, (e, Vertex))]
forall a b. (a -> b) -> [a] -> [b]
map Edge e -> (Vertex, (e, Vertex))
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) <- Graph e -> [Edge e]
forall e. Graph e -> [Edge e]
edges Graph e
g ]

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

vertices :: Graph e -> [Vertex]
vertices :: forall e. Graph e -> [Vertex]
vertices = Array Vertex [(e, Vertex)] -> [Vertex]
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 <- Graph e -> [Vertex]
forall e. Graph e -> [Vertex]
vertices Graph e
g, (e
e, Vertex
w) <- Graph e
gGraph e -> Vertex -> [(e, Vertex)]
forall 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 = ([(e, Vertex)] -> [Vertex]) -> Array Vertex [(e, Vertex)] -> Graph
forall a b. (a -> b) -> Array Vertex a -> Array Vertex b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((e, Vertex) -> Vertex) -> [(e, Vertex)] -> [Vertex]
forall a b. (a -> b) -> [a] -> [b]
map (e, Vertex) -> Vertex
forall a b. (a, b) -> b
snd)

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