{-# LANGUAGE DoAndIfThenElse #-} -- | Implementation of a graph with each node identified by a unique key. -- It is a provisional module and it might be replaced by the standard -- graph from containers package in the future. module Data.Named.Graph ( Graph (..) , mkGraph , node , edges , roots , disjointForest ) where import qualified Data.Set as S import qualified Data.Map as M import qualified Data.Tree as T import Data.Named.Tree -- | A graph. data Graph k v = Graph { nodeMap :: M.Map k v , edgeMap :: M.Map k [k] } -- | Make a graph from a list of (key, value, [children keys]) tuples. mkGraph :: Ord k => [(k, v, [k])] -> Graph k v mkGraph xs = Graph ns es where ns = M.fromList [(k, v) | (k, v, _) <- xs] es = M.fromList [(k, ks) | (k, _, ks) <- xs] -- | Get node with the given key. node :: (Show k, Ord k) => Graph k v -> k -> v node g k = case M.lookup k (nodeMap g) of Nothing -> error $ "node: key " ++ show k ++ " not in the nodes map" Just v -> v {-# INLINE node #-} -- | Get keys of adjacent nodes for the given node key. edges :: (Show k, Ord k) => Graph k v -> k -> [k] edges g k = case M.lookup k (edgeMap g) of Nothing -> error $ "edges: key " ++ show k ++ " not in the edges map" Just v -> v {-# INLINE edges #-} -- | Return all graph roots (i.e. nodes with no parents). roots :: Ord k => Graph k v -> [k] roots g = let desc = S.fromList . concat . M.elems $ edgeMap g in [k | (k, _) <- M.assocs (nodeMap g), not (k `S.member` desc)] generate :: (Show k, Ord k) => Graph k v -> k -> T.Tree k generate g k = T.Node k (map (generate g) (edges g k)) -- | A stateful monad for forest pruning. newtype RanM a = RanM { runRanM :: Int -> (a, Int) } instance Monad RanM where return x = RanM $ \s -> (x, s) RanM v >>= f = RanM $ \s -> case v s of (x, s') -> runRanM (f x) s' run :: RanM a -> a run act = fst (runRanM act (-1)) contains :: Int -> RanM Bool contains k = RanM $ \m -> (k <= m, m) include :: Int -> RanM () include k = RanM $ \_ -> ((), k) chop :: T.Forest (k, Span) -> RanM (T.Forest (k, Span)) chop [] = return [] chop (T.Node (k, s) ts : us) = do visited <- contains (end s) if visited then chop us else do as <- chop ts include (end s) bs <- chop us return (T.Node (k, s) as : bs) prune :: (k -> Int) -> T.Forest k -> T.Forest k prune f = unSpanForest . run . chop . sortForest . spanForest f -- | Spanning-like forest of a DAG. Trees in the resulting forest are -- disjoint with respect to their ranges. It is not checked if the input -- graph is actually a DAG. disjointForest :: (Show k, Ord k) => (k -> Int) -> Graph k v -> T.Forest k disjointForest f g = prune f . map (generate g) . roots $ g