module RegAlloc.Data.Graphs (
    -- * Graphs
    -- ** Type
    Graph,

    -- ** Construction
    graph,
    unconnected,

    -- ** Properties
    -- *** Simple
    vertices,
    edges,
    successors,
    predecessors,
    outDegree,
    inDegree,

    -- *** Advanced
    spanningTree,
    wccs,

    -- ** Operations
    transpose,
    edgeUnion,
    edgeDiff,
    edgeMap,
    edgeFilter,
    concatenate,
    closure,
    undirected,

    -- ** Utilities
    ifVertexSetsAreEqual,

    -- * Decompositions
    Decomposition,
    ComponentID,
    componentID
) where
    import Data.Array
    import Data.Set
    import Control.Arrow
    import Control.Monad.Reader
    import Control.Monad.State
    import RegAlloc.Data.IndexedSets
    import RegAlloc.Data.Trees

    -- * Graphs
    -- ** Type
    data Ord vertex => Graph vertex = Graph (IndexedSet vertex)
                                            (Array Int (Set Int))
                                            (Array Int (Set Int))
#ifndef __HADDOCK__
                       deriving Eq
#else
    instance Ord vertex => Eq (Graph vertex)
#endif

    -- ** Construction
    graph :: Ord vertex => Set vertex -> Set (vertex,vertex) -> Graph vertex
    graph vertices edges = graph' (indexedSet vertices) edges

    graph' :: Ord vertex => IndexedSet vertex -> Set (vertex,vertex) -> Graph vertex
    graph' indexedVertices edges
        = graph'' indexedVertices (map (join (***) (indexOf indexedVertices)) (setToList edges))

    graph'' :: Ord vertex => IndexedSet vertex -> [(Int,Int)] -> Graph vertex
    graph'' indexedVertices indexPairs
        = Graph indexedVertices
                (neighborTable indexPairs)
                (neighborTable (map (\(index1,index2) -> (index2,index1)) indexPairs))
        where
            neighborTable :: [(Int,Int)] -> Array Int (Set Int)
            neighborTable = accumArray addToSet emptySet (1,card indexedVertices)

    unconnected :: Ord vertex => Set vertex -> Graph vertex
    unconnected vertices = join (Graph (indexedSet vertices)) $
                           (listArray (1,cardinality vertices) (repeat emptySet))

    -- ** Properties
    -- *** Simple
    vertices :: Ord vertex => Graph vertex -> Set vertex
    vertices (Graph indexedVertices _ _) = elements indexedVertices

    edges :: Ord vertex => Graph vertex -> Set (vertex,vertex)
    edges (Graph indexedVertices forwardTable _)
        = mkSet [join (***) (elementAt indexedVertices) (startIndex,forwardIndex) |
                 (startIndex,forwardIndices) <- assocs forwardTable,
                 forwardIndex                <- setToList forwardIndices]

    successors :: Ord vertex => vertex -> Graph vertex -> Set vertex
    successors vertex (Graph indexedVertices forwardTable _)
        = mapSet (elementAt indexedVertices)
                 (forwardTable ! indexOf indexedVertices vertex)

    predecessors :: Ord vertex => vertex -> Graph vertex -> Set vertex
    predecessors vertex = successors vertex . transpose

    outDegree :: Ord vertex => vertex -> Graph vertex -> Int
    outDegree vertex (Graph indexedVertices forwardTable _)
        = cardinality (forwardTable ! indexOf indexedVertices vertex)

    inDegree :: Ord vertex => vertex -> Graph vertex -> Int
    inDegree vertex = outDegree vertex . transpose

    -- *** Advanced
    spanningTree :: Ord vertex => vertex -> Graph vertex -> Tree vertex
    spanningTree root (Graph indexedVertices forwardTable _)
        = fmap (elementAt indexedVertices)
               (spanningIndexTree (indexOf indexedVertices root) forwardTable)

    spanningIndexTree :: Int -> Array Int (Set Int) -> Tree Int
    spanningIndexTree rootIndex forwardTable
        = evalState (findTree rootIndex) (mkSet (indices forwardTable))
        where
            findTree rootIndex
                = do
                      modify (`delFromSet` rootIndex)
                      rootChildIndices
                          <- gets (\toDo -> [childIndex |
                                             childIndex <- setToList (forwardTable ! rootIndex),
                                             not (childIndex `elementOf` toDo)])
                          -- Be careful.  Set intersection would result in O(n?) time for the whole
                          -- algorithm.
                      subtrees <- sequence [findTree rootChildIndex |
                                            rootChildIndex <- rootChildIndices]
                      return (tree rootIndex (forest subtrees))

    wccs :: Ord vertex => Graph vertex -> Decomposition vertex
    wccs graph@(Graph indexedVertices _ _)
        = let
              indexWCCs = findIndexWCCs (mkSet [1..card indexedVertices])
          in
          Decomposition indexedVertices
                        (array (1,length indexWCCs)
                               [(vertexIndex,componentIndex) |
                                (vertexIndices,componentIndex) <- zip indexWCCs [1..],
                                vertexIndex <- vertexIndices])
        where
            findIndexWCCs :: Set Int -> [[Int]]
            findIndexWCCs toDo
                | isEmptySet toDo
                    = []
                | otherwise
                    = let
                          component = flattenTree (spanningIndexTree (head (setToList toDo))
                                                                     undirectedTable)
                      in
                      component : findIndexWCCs (foldr (flip delFromSet) toDo component)

            Graph _ undirectedTable _ = undirected graph

    -- ** Operations
    transpose :: Ord vertex => Graph vertex -> Graph vertex
    transpose (Graph indexedVertices forwardTable backwardTable)
        = Graph indexedVertices backwardTable forwardTable

    edgeUnion :: Ord vertex => Graph vertex -> Graph vertex -> Graph vertex
    edgeUnion (Graph indexedVertices1 forwardTable1 backwardTable1)
              (Graph indexedVertices2 forwardTable2 backwardTable2)
        = ifVertexSetsAreEqual indexedVertices1
                               indexedVertices2
                               (Graph indexedVertices1
                                      (unionTable forwardTable1 forwardTable2)
                                      (unionTable backwardTable1 backwardTable2))
        where
            unionTable :: Array Int (Set Int) -> Array Int (Set Int) -> Array Int (Set Int)
            unionTable table1 table2 = listArray (bounds table1)
                                                 (zipWith union (elems table1) (elems table2))

    edgeDiff :: Ord vertex => Graph vertex -> Graph vertex -> Graph vertex
    edgeDiff (Graph indexedVertices1 forwardTable1 backwardTable1)
             (Graph indexedVertices2 forwardTable2 backwardTable2)
        = ifVertexSetsAreEqual indexedVertices1
                         indexedVertices2
                         (Graph indexedVertices1
                                (diffTable forwardTable1 forwardTable2)
                                (diffTable backwardTable2 backwardTable2))
        where
            diffTable :: Array Int (Set Int) -> Array Int (Set Int) -> Array Int (Set Int)
            diffTable table1 table2 = listArray (bounds table1)
                                                (zipWith minusSet (elems table1) (elems table2))

    edgeMap :: Ord vertex => ((vertex,vertex) -> (vertex,vertex)) -> Graph vertex -> Graph vertex
    edgeMap mapping graph@(Graph indexedVertices _ _)
        = graph' indexedVertices (mapSet mapping (edges graph))

    edgeFilter :: Ord vertex => ((vertex,vertex) -> Bool) -> Graph vertex -> Graph vertex
    edgeFilter predicate graph@(Graph indexedVertices _ _)
        = graph' indexedVertices (mkSet (filter predicate (setToList (edges graph))))

    concatenate :: Ord vertex => Graph vertex -> Graph vertex -> Graph vertex
    concatenate (Graph indexedVertices1 _ backwardTable1)
                (Graph indexedVertices2 forwardTable2 _)
        = ifVertexSetsAreEqual indexedVertices1
                               indexedVertices2
                               (graph'' indexedVertices1
                                        (concatMap
                                             (uncurry (liftM2 (,)) . join (***) setToList)
                                             (zip (elems backwardTable1) (elems forwardTable2))))

    closure :: Ord vertex => Graph vertex -> Graph vertex
    closure (Graph indexedVertices forwardTable _)
        = graph'' indexedVertices
                  [(index1,index2) | index1 <- [1..card indexedVertices],
                                     index2 <- flattenTree (spanningIndexTree index1 forwardTable)]

    undirected :: Ord vertex => Graph vertex -> Graph vertex
    undirected graph = graph `edgeUnion` transpose graph

    -- ** Utilities
    ifVertexSetsAreEqual :: Ord vertex => IndexedSet vertex -> IndexedSet vertex -> value -> value
    ifVertexSetsAreEqual indexedVertices1 indexedVertices2
        | indexedVertices1 == indexedVertices2
            = id
        | otherwise
            = error "RegAlloc.Data.Graphs.ifVertexSetsAreEqual: vertex sets don't match"

    -- * Decompositions
    data Ord element => Decomposition element = Decomposition (IndexedSet element) (Array Int Int)
#ifndef __HADDOCK__
                        deriving Eq
#else
    instance Ord element => Eq (Decomposition element)
#endif

    newtype Ord element => ComponentID element = ComponentID Int
#ifndef __HADDOCK__
                           deriving (Eq, Ord)
#else
    instance Ord element => Eq (ComponentID element)
    instance Ord element => Ord (ComponentID element)
#endif

    componentCount :: Ord element => Decomposition element -> Int
    componentCount (Decomposition _ mapping) = snd (bounds mapping)

    componentID :: Ord element => Decomposition element -> element -> ComponentID element
    componentID (Decomposition indexedElements mapping) element
        = ComponentID (mapping ! indexOf indexedElements element)