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

    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

    import qualified Data.Map.Strict as Map
    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
        foldl' f z = Foldable.foldl' f z . graphMap
        foldr' f z = Foldable.foldr' f z . graphMap
    #ifdef MIN_VERSION_base
    #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)