Proposed new graph interface (refactor of PackageIndex)

[ezyang]

I propose we introduce a general, user-friendly graph type to Cabal (Data.Graph is NOT user-friendly), which supports arbitrary key types (as opposed to PackageIndex, which bakes in UnitId as the key.) This will give us more flexibility to pick different keys for solver plans, install plans, etc.

Here is the proposed module:

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Graph.Node
-- Copyright   :  (c) Edward Z. Yang
-- License     :  BSD3
--
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  portable
--
-- This alternative interface to 'Data.Graph' stores nodes (identified
-- by the 'IsNode' type class) which are associated with a key and
-- record the keys of their neighbors.
--
-- Special attention is paid to graphs which are "broken"; i.e., they
-- have neighbor keys which are not in the map.
--
-----------------------------------------------------------------------------

module Data.Graph.Node (SCC(..), module Data.Graph.Node) where

import Prelude hiding (lookup)
import Data.Graph (SCC(..))
import qualified Data.Graph as G
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Array as Array
import Data.Array ((!))
import Data.List hiding (lookup)
import Data.Ord
import qualified Data.Tree as Tree
import Data.Either

-- | 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])]
    }

class Ord (Key a) => IsNode a where
    type Key a :: *
    nodeKey :: a -> Key a
    nodeNeighbors :: a -> [Key a]

fromMap :: IsNode a => Map (Key a) a -> Graph a
fromMap m
    = Graph { graphMap = m
            , graphForward = g
            , graphAdjoint = G.transposeG g
            , graphVertexToNode = vertex_to_node
            , graphKeyToVertex = key_to_vertex
            , graphBroken = broken
            }
  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 ]
    broken = filter (not . null . snd) (zip ns brokenEdges)

    g = Array.listArray bounds edges

    ns              = sortBy (comparing nodeKey) (Map.elems m)
    vertices        = zip (map nodeKey ns) [0..]
    vertex_map      = Map.fromList 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)

fromList :: IsNode a => [a] -> Graph a
fromList ns = fromMap
            . Map.fromList
            . map (\n -> (nodeKey n, n))
            $ ns

-- Map-like operations

-- TODO: More efficient graph rebuilding strategy

insert :: IsNode a => a -> Graph a -> Graph a
insert n g = fromMap (Map.insert (nodeKey n) n (toMap g))

deleteKey :: IsNode a => Key a -> Graph a -> Graph a
deleteKey k g = fromMap (Map.delete k (toMap g))

lookup :: IsNode a => Key a -> Graph a -> Maybe a
lookup k g = Map.lookup k (toMap g)

toList :: Graph a -> [a]
toList g = Map.elems (toMap g)

toMap :: Graph a -> Map (Key a) a
toMap = graphMap

-- Graph-like operations

-- From stronglyConnCompR in Data.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)

cycles :: Graph a -> [[a]]
cycles g = [ vs | CyclicSCC vs <- stronglyConnComp g ]

broken :: IsNode a => Graph a -> [(a, [Key a])]
broken g = graphBroken g

closure :: IsNode a => Graph a -> [Key a] -> Maybe (Graph a)
closure g ks = do
    vs <- mapM (graphKeyToVertex g) ks
    return (fromList (decodeVertexForest g (G.dff (graphForward g))))

revClosure :: IsNode a => Graph a -> [Key a] -> Maybe [a]
revClosure g ks = do
    vs <- mapM (graphKeyToVertex g) ks
    return (decodeVertexForest g (G.dff (graphAdjoint g)))

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

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

topSort :: Graph a -> [a]
topSort g = map (graphVertexToNode g) $ G.topSort (graphForward g)

revTopSort :: Graph a -> [a]
revTopSort g = map (graphVertexToNode g) $ G.topSort (graphAdjoint g)

CC @dcoutts