{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE BangPatterns #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Distribution.Compat.Graph
-- Copyright   :  (c) Edward Z. Yang 2016
-- License     :  BSD3
--
-- Maintainer  :  [email protected]
-- Stability   :  experimental
-- Portability :  portable
--
-- A data type representing directed graphs, backed by "Data.Graph".
-- It is strict in the node type.
--
-- This is an alternative interface to "Data.Graph".  In this interface,
-- nodes (identified by the 'IsNode' type class) are associated with a
-- key and record the keys of their neighbors.  This interface is more
-- convenient than 'Data.Graph.Graph', which requires vertices to be
-- explicitly handled by integer indexes.
--
-- The current implementation has somewhat peculiar performance
-- characteristics.  The asymptotics of all map-like operations mirror
-- their counterparts in "Data.Map".  However, to perform a graph
-- operation, we first must build the "Data.Graph" representation, an
-- operation that takes /O(V + E log V)/.  However, this operation can
-- be amortized across all queries on that particular graph.
--
-- Some nodes may be broken, i.e., refer to neighbors which are not
-- stored in the graph.  In our graph algorithms, we transparently
-- ignore such edges; however, you can easily query for the broken
-- vertices of a graph using 'broken' (and should, e.g., to ensure that
-- a closure of a graph is well-formed.)  It's possible to take a closed
-- subset of a broken graph and get a well-formed graph.
--
-----------------------------------------------------------------------------

module Distribution.Compat.Graph (
    -- * Graph type
    Graph,
    IsNode(..),
    -- * Query
    null,
    size,
    member,
    lookup,
    -- * Construction
    empty,
    insert,
    deleteKey,
    deleteLookup,
    -- * Combine
    unionLeft,
    unionRight,
    -- * Graph algorithms
    stronglyConnComp,
    SCC(..),
    cycles,
    broken,
    neighbors,
    revNeighbors,
    closure,
    revClosure,
    topSort,
    revTopSort,
    -- * Conversions
    -- ** Maps
    toMap,
    -- ** Lists
    fromDistinctList,
    toList,
    keys,
    -- ** Sets
    keysSet,
    -- ** Graphs
    toGraph,
    -- * Node type
    Node(..),
    nodeValue,
) where

-- For bootstrapping GHC
#ifdef MIN_VERSION_containers
#if MIN_VERSION_containers(0,5,0)
#define HAVE_containers_050
#endif
#endif

import Prelude ()
import qualified Distribution.Compat.Prelude as Prelude
import Distribution.Compat.Prelude hiding (lookup, null, empty)

import Data.Graph (SCC(..))
import qualified Data.Graph as G

#ifdef HAVE_containers_050
import qualified Data.Map.Strict as Map
#else
import qualified Data.Map as Map
#endif
import qualified Data.Set as Set
import qualified Data.Array as Array
import Data.Array ((!))
import qualified Data.Tree as Tree
import Data.Either (partitionEithers)
import qualified Data.Foldable as Foldable

-- | A graph of nodes @a@.  The nodes are expected to have instance
-- of class 'IsNode'.
data Graph a
    = Graph {
        graphMap          :: !(Map (Key a) a),
        -- Lazily cached graph representation
        graphForward      :: G.Graph,
        graphAdjoint      :: G.Graph,
        graphVertexToNode :: G.Vertex -> a,
        graphKeyToVertex  :: Key a -> Maybe G.Vertex,
        graphBroken       :: [(a, [Key a])]
    }
    deriving (Typeable)

-- NB: Not a Functor! (or Traversable), because you need
-- to restrict Key a ~ Key b.  We provide our own mapping
-- functions.

-- General strategy is most operations are deferred to the
-- Map representation.

instance Show a => Show (Graph a) where
    show = show . toList

instance (IsNode a, Read a, Show (Key a)) => Read (Graph a) where
    readsPrec d s = map (\(a,r) -> (fromDistinctList a, r)) (readsPrec d s)

instance (IsNode a, Binary a, Show (Key a)) => Binary (Graph a) where
    put x = put (toList x)
    get = fmap fromDistinctList get

instance (Eq (Key a), Eq a) => Eq (Graph a) where
    g1 == g2 = graphMap g1 == graphMap g2

instance Foldable.Foldable Graph where
    fold = Foldable.fold . graphMap
    foldr f z = Foldable.foldr f z . graphMap
    foldl f z = Foldable.foldl f z . graphMap
    foldMap f = Foldable.foldMap f . graphMap
#ifdef MIN_VERSION_base
#if MIN_VERSION_base(4,6,0)
    foldl' f z = Foldable.foldl' f z . graphMap
    foldr' f z = Foldable.foldr' f z . graphMap
#endif
#if MIN_VERSION_base(4,8,0)
    length = Foldable.length . graphMap
    null   = Foldable.null   . graphMap
    toList = Foldable.toList . graphMap
    elem x = Foldable.elem x . graphMap
    maximum = Foldable.maximum . graphMap
    minimum = Foldable.minimum . graphMap
    sum     = Foldable.sum     . graphMap
    product = Foldable.product . graphMap
#endif
#endif

instance (NFData a, NFData (Key a)) => NFData (Graph a) where
    rnf Graph {
        graphMap = m,
        graphForward = gf,
        graphAdjoint = ga,
        graphVertexToNode = vtn,
        graphKeyToVertex = ktv,
        graphBroken = b
    } = gf `seq` ga `seq` vtn `seq` ktv `seq` b `seq` rnf m

-- TODO: Data instance?

-- | The 'IsNode' class is used for datatypes which represent directed
-- graph nodes.  A node of type @a@ is associated with some unique key of
-- type @'Key' a@; given a node we can determine its key ('nodeKey')
-- and the keys of its neighbors ('nodeNeighbors').
class Ord (Key a) => IsNode a where
    type Key a :: *
    nodeKey :: a -> Key a
    nodeNeighbors :: a -> [Key a]

instance (IsNode a, IsNode b, Key a ~ Key b) => IsNode (Either a b) where
    type Key (Either a b) = Key a
    nodeKey (Left x)  = nodeKey x
    nodeKey (Right x) = nodeKey x
    nodeNeighbors (Left x)  = nodeNeighbors x
    nodeNeighbors (Right x) = nodeNeighbors x

-- | A simple, trivial data type which admits an 'IsNode' instance.
data Node k a = N a k [k]
    deriving (Show, Eq)

-- | Get the value from a 'Node'.
nodeValue :: Node k a -> a
nodeValue (N a _ _) = a

instance Functor (Node k) where
    fmap f (N a k ks) = N (f a) k ks

instance Ord k => IsNode (Node k a) where
    type Key (Node k a) = k
    nodeKey (N _ k _) = k
    nodeNeighbors (N _ _ ks) = ks

-- TODO: Maybe introduce a typeclass for items which just
-- keys (so, Key associated type, and nodeKey method).  But
-- I didn't need it here, so I didn't introduce it.

-- Query

-- | /O(1)/. Is the graph empty?
null :: Graph a -> Bool
null = Map.null . toMap

-- | /O(1)/. The number of nodes in the graph.
size :: Graph a -> Int
size = Map.size . toMap

-- | /O(log V)/. Check if the key is in the graph.
member :: IsNode a => Key a -> Graph a -> Bool
member k g = Map.member k (toMap g)

-- | /O(log V)/. Lookup the node at a key in the graph.
lookup :: IsNode a => Key a -> Graph a -> Maybe a
lookup k g = Map.lookup k (toMap g)

-- Construction

-- | /O(1)/. The empty graph.
empty :: IsNode a => Graph a
empty = fromMap Map.empty

-- | /O(log V)/. Insert a node into a graph.
insert :: IsNode a => a -> Graph a -> Graph a
insert !n g = fromMap (Map.insert (nodeKey n) n (toMap g))

-- | /O(log V)/. Delete the node at a key from the graph.
deleteKey :: IsNode a => Key a -> Graph a -> Graph a
deleteKey k g = fromMap (Map.delete k (toMap g))

-- | /O(log V)/. Lookup and delete.  This function returns the deleted
-- value if it existed.
deleteLookup :: IsNode a => Key a -> Graph a -> (Maybe a, Graph a)
deleteLookup k g =
    let (r, m') = Map.updateLookupWithKey (\_ _ -> Nothing) k (toMap g)
    in (r, fromMap m')

-- Combining

-- | /O(V + V')/. Right-biased union, preferring entries
-- from the second map when conflicts occur.
-- @'nodeKey' x = 'nodeKey' (f x)@.
unionRight :: IsNode a => Graph a -> Graph a -> Graph a
unionRight g g' = fromMap (Map.union (toMap g') (toMap g))

-- | /O(V + V')/. Left-biased union, preferring entries from
-- the first map when conflicts occur.
unionLeft :: IsNode a => Graph a -> Graph a -> Graph a
unionLeft = flip unionRight

-- Graph-like operations

-- | /Ω(V + E)/. Compute the strongly connected components of a graph.
-- Requires amortized construction of graph.
stronglyConnComp :: Graph a -> [SCC a]
stronglyConnComp g = map decode forest
  where
    forest = G.scc (graphForward g)
    decode (Tree.Node v [])
        | mentions_itself v = CyclicSCC  [graphVertexToNode g v]
        | otherwise         = AcyclicSCC (graphVertexToNode g v)
    decode other = CyclicSCC (dec other [])
        where dec (Tree.Node v ts) vs
                = graphVertexToNode g v : foldr dec vs ts
    mentions_itself v = v `elem` (graphForward g ! v)
-- Implementation copied from 'stronglyConnCompR' in 'Data.Graph'.

-- | /Ω(V + E)/. Compute the cycles of a graph.
-- Requires amortized construction of graph.
cycles :: Graph a -> [[a]]
cycles g = [ vs | CyclicSCC vs <- stronglyConnComp g ]

-- | /O(1)/.  Return a list of nodes paired with their broken
-- neighbors (i.e., neighbor keys which are not in the graph).
-- Requires amortized construction of graph.
broken :: Graph a -> [(a, [Key a])]
broken g = graphBroken g

-- | Lookup the immediate neighbors from a key in the graph.
-- Requires amortized construction of graph.
neighbors :: Graph a -> Key a -> Maybe [a]
neighbors g k = do
    v <- graphKeyToVertex g k
    return (map (graphVertexToNode g) (graphForward g ! v))

-- | Lookup the immediate reverse neighbors from a key in the graph.
-- Requires amortized construction of graph.
revNeighbors :: Graph a -> Key a -> Maybe [a]
revNeighbors g k = do
    v <- graphKeyToVertex g k
    return (map (graphVertexToNode g) (graphAdjoint g ! v))

-- | Compute the subgraph which is the closure of some set of keys.
-- Returns @Nothing@ if one (or more) keys are not present in
-- the graph.
-- Requires amortized construction of graph.
closure :: Graph a -> [Key a] -> Maybe [a]
closure g ks = do
    vs <- traverse (graphKeyToVertex g) ks
    return (decodeVertexForest g (G.dfs (graphForward g) vs))

-- | Compute the reverse closure of a graph from some set
-- of keys.  Returns @Nothing@ if one (or more) keys are not present in
-- the graph.
-- Requires amortized construction of graph.
revClosure :: Graph a -> [Key a] -> Maybe [a]
revClosure g ks = do
    vs <- traverse (graphKeyToVertex g) ks
    return (decodeVertexForest g (G.dfs (graphAdjoint g) vs))

flattenForest :: Tree.Forest a -> [a]
flattenForest = concatMap Tree.flatten

decodeVertexForest :: Graph a -> Tree.Forest G.Vertex -> [a]
decodeVertexForest g = map (graphVertexToNode g) . flattenForest

-- | Topologically sort the nodes of a graph.
-- Requires amortized construction of graph.
topSort :: Graph a -> [a]
topSort g = map (graphVertexToNode g) $ G.topSort (graphForward g)

-- | Reverse topologically sort the nodes of a graph.
-- Requires amortized construction of graph.
revTopSort :: Graph a -> [a]
revTopSort g = map (graphVertexToNode g) $ G.topSort (graphAdjoint g)

-- Conversions

-- | /O(1)/. Convert a map from keys to nodes into a graph.
-- The map must satisfy the invariant that
-- @'fromMap' m == 'fromList' ('Data.Map.elems' m)@;
-- if you can't fulfill this invariant use @'fromList' ('Data.Map.elems' m)@
-- instead.  The values of the map are assumed to already
-- be in WHNF.
fromMap :: IsNode a => Map (Key a) a -> Graph a
fromMap m
    = Graph { graphMap = m
            -- These are lazily computed!
            , graphForward = g
            , graphAdjoint = G.transposeG g
            , graphVertexToNode = vertex_to_node
            , graphKeyToVertex = key_to_vertex
            , graphBroken = broke
            }
  where
    try_key_to_vertex k = maybe (Left k) Right (key_to_vertex k)

    (brokenEdges, edges)
        = unzip
        $ [ partitionEithers (map try_key_to_vertex (nodeNeighbors n))
          | n <- ns ]
    broke = filter (not . Prelude.null . snd) (zip ns brokenEdges)

    g = Array.listArray bounds edges

    ns              = Map.elems m -- sorted ascending
    vertices        = zip (map nodeKey ns) [0..]
    vertex_map      = Map.fromAscList vertices
    key_to_vertex k = Map.lookup k vertex_map

    vertex_to_node vertex = nodeTable ! vertex

    nodeTable   = Array.listArray bounds ns
    bounds = (0, Map.size m - 1)

-- | /O(V log V)/. Convert a list of nodes (with distinct keys) into a graph.
fromDistinctList :: (IsNode a, Show (Key a)) => [a] -> Graph a
fromDistinctList = fromMap
                 . Map.fromListWith (\_ -> duplicateError)
                 . map (\n -> n `seq` (nodeKey n, n))
  where
    duplicateError n = error $ "Graph.fromDistinctList: duplicate key: "
                            ++ show (nodeKey n)

-- Map-like operations

-- | /O(V)/. Convert a graph into a list of nodes.
toList :: Graph a -> [a]
toList g = Map.elems (toMap g)

-- | /O(V)/. Convert a graph into a list of keys.
keys :: Graph a -> [Key a]
keys g = Map.keys (toMap g)

-- | /O(V)/. Convert a graph into a set of keys.
keysSet :: Graph a -> Set.Set (Key a)
keysSet g = Map.keysSet (toMap g)

-- | /O(1)/. Convert a graph into a map from keys to nodes.
-- The resulting map @m@ is guaranteed to have the property that
-- @'Prelude.all' (\(k,n) -> k == 'nodeKey' n) ('Data.Map.toList' m)@.
toMap :: Graph a -> Map (Key a) a
toMap = graphMap

-- Graph-like operations

-- | /O(1)/. Convert a graph into a 'Data.Graph.Graph'.
-- Requires amortized construction of graph.
toGraph :: Graph a -> (G.Graph, G.Vertex -> a, Key a -> Maybe G.Vertex)
toGraph g = (graphForward g, graphVertexToNode g, graphKeyToVertex g)