Distribution/Compat/Graph.hs

    {-# 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  :  cabal-dev@haskell.org
    -- 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)