-
-
Save alpmestan/a25e38107d41abefbb60c9953957e01b to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| {-# LANGUAGE RankNTypes #-} | |
| {-# LANGUAGE RecursiveDo #-} | |
| {-# LANGUAGE DeriveFunctor #-} | |
| {-# LANGUAGE GeneralizedNewtypeDeriving #-} | |
| -- | A simple graph library | |
| module Graph | |
| ( -- * Creating graphs | |
| runGraph | |
| , newGraph | |
| , node | |
| , Graph(..) | |
| , Vertex(..) | |
| , Node(..) | |
| -- * Topological sort | |
| , topologicalSort | |
| -- * Drawing graphs | |
| , viewGraph | |
| , graphviz | |
| -- * Internal types | |
| , MGraph(..) | |
| , Vertex_(..) | |
| ) where | |
| import Control.Monad | |
| import Control.Monad.ST | |
| import Data.List | |
| import Data.STRef | |
| import Data.Vector (Vector) | |
| import Data.Vector.Mutable (MVector) | |
| import qualified Data.Vector as V | |
| import qualified Data.Vector.Mutable as MV | |
| import System.IO.Temp | |
| import System.Process | |
| -- * Graph API | |
| -- | a 'Node' is just an integer identifier. | |
| newtype Node = Node | |
| { nodeId :: Int | |
| } deriving (Eq, Show, Ord, Num) | |
| -- | An "incomplete" version of 'Vertex' where we do not | |
| -- yet have the output nodes handy, nor the labels. Used | |
| -- when building the graph. | |
| data Vertex_ a = Vertex_ | |
| { vlabel :: !a | |
| , vid :: Node | |
| , vinputs :: [Node] | |
| } deriving (Functor, Eq, Show) | |
| -- | A "mutable graph". | |
| data MGraph s a = MGraph | |
| { graphData :: !(STRef s (MVector s (Vertex_ a))) | |
| , nextNode :: {-# UNPACK #-} !(STRef s Node) | |
| } | |
| -- | A vertex in our 'Graph's. Contains the 'label', | |
| -- the node 'ident'ifier, the 'inputs' nodes and 'outputs' | |
| -- nodes of this vertex. | |
| data Vertex a = Vertex | |
| { label :: !a | |
| , ident :: Node | |
| , inputs :: [(Node, a)] | |
| , outputs :: [(Node, a)] | |
| } deriving (Functor, Eq, Show) | |
| -- | A 'Graph' as collection of vertices. | |
| newtype Graph a = Graph | |
| { graphArray :: Vector (Vertex a) | |
| } deriving (Functor, Eq, Show) | |
| -- | Initialize a new mutable graph. | |
| newGraph :: ST s (MGraph s a) | |
| newGraph = do | |
| mv0 <- MV.new 8 | |
| v0 <- newSTRef mv0 | |
| n0 <- newSTRef 0 | |
| return (MGraph v0 n0) | |
| -- | @node g a xs@ adds a vertex with label @a@ | |
| -- and inputs @xs@ to the graph @g@, returning | |
| -- the identifier of the newly added node. | |
| node :: MGraph s a -> a -> [Node] -> ST s Node | |
| node mg a xs = do | |
| mv <- readSTRef (graphData mg) | |
| Node nodeid <- readSTRef (nextNode mg) | |
| let arrSize = MV.length mv | |
| available = arrSize - nodeid | |
| getTarget | |
| | available < 1 = do mv' <- MV.grow mv (2 * arrSize) | |
| writeSTRef (graphData mg) mv' | |
| return mv' | |
| | otherwise = return mv | |
| targetVec <- getTarget | |
| MV.write targetVec nodeid (Vertex_ a (Node nodeid) xs) | |
| modifySTRef' (nextNode mg) (+1) | |
| return (Node nodeid) | |
| buildOutputs :: Vector (Vertex_ a) -> Vector [(Node, a)] | |
| buildOutputs v = runST $ do | |
| mv <- MV.new (V.length v) | |
| forM_ v $ \(Vertex_ a n@(Node i) is) -> do | |
| MV.write mv i [] | |
| forM_ is $ \(Node j) -> | |
| MV.modify mv (++[(n, a)]) j | |
| V.freeze mv | |
| buildGraph :: Vector (Vertex_ a) -> Graph a | |
| buildGraph v = Graph (V.map f v) | |
| where f (Vertex_ a nid@(Node i) is) = Vertex a nid (map lkp is) (outputs V.! i) | |
| lkp n@(Node i) = (n, vlabel (v V.! i)) | |
| outputs = buildOutputs v | |
| freezeGraph :: MGraph s a -> ST s (Graph a) | |
| freezeGraph mg = do | |
| numNodes <- (\(Node n) -> n) <$> readSTRef (nextNode mg) | |
| mv <- MV.take numNodes <$> readSTRef (graphData mg) | |
| v <- V.freeze mv | |
| return (buildGraph v) | |
| -- | Turn a computation returning a mutable graph into | |
| -- an immutable 'Graph'. | |
| -- | |
| -- @ | |
| -- ex = runGraph $ do | |
| -- g <- newGraph | |
| -- a <- node g \"A\" [] | |
| -- b <- node g \"B\" [a] | |
| -- c <- node g \"C\" [a, b] | |
| -- return g | |
| -- @ | |
| runGraph :: (forall s. ST s (MGraph s a)) -> Graph a | |
| runGraph f = runST (freezeGraph =<< f) | |
| buildEdges :: Vector (Vertex a) -> [(Node, Node)] | |
| buildEdges v = [ (ni, n) | |
| | Vertex _ n is _ <- V.toList v | |
| , (ni, _) <- is | |
| ] | |
| -- | Returns an error if the graph has cycles, a list of | |
| -- all the vertices of the graph sorted in topological order | |
| -- otherwise. | |
| topologicalSort :: Graph a -> [Vertex a] | |
| topologicalSort (Graph v) = | |
| let (nodes, result) = runST $ do | |
| nodes0 <- newSTRef startNodes | |
| edges0 <- newSTRef edgeList | |
| res0 <- newSTRef [] | |
| go nodes0 edges0 res0 | |
| in | |
| if not (null nodes) | |
| then error "topologicalSort: the graph has cycles" | |
| else result | |
| where startNodes = V.toList (V.filter (null . inputs) v) | |
| edgeList = buildEdges v | |
| go nodes edges res = do | |
| ns <- readSTRef nodes | |
| case ns of | |
| [] -> (,) <$> readSTRef edges <*> readSTRef res | |
| (vtx:vs) -> do | |
| modifySTRef' nodes tail | |
| modifySTRef' res (++[vtx]) | |
| forM_ (outputs $ v V.! nodeId (ident vtx)) $ \(i, a) -> do | |
| es <- readSTRef edges | |
| writeSTRef edges (es \\ [(ident vtx, i)]) | |
| let noOtherIncomingEdge = (==1) . length $ filter ((==i) . snd) es | |
| when noOtherIncomingEdge $ | |
| modifySTRef' nodes (\xs -> xs ++ [v V.! nodeId i]) | |
| go nodes edges res | |
| -- * Graphviz | |
| -- | Generates the graphviz description of the graph | |
| graphviz :: Graph String -> String | |
| graphviz g = unlines . wrapGraph . foldMap collect $ graphArray g | |
| where collect (Vertex a _ is os) = (" \"" ++ a ++ "\";") : map (ppEdgeToFrom a . snd) is | |
| ppEdgeToFrom lbl1 lbl2 = " \"" ++ lbl2 ++ "\" -> \"" ++ lbl1 ++ "\";" | |
| wrapGraph xs = [ "digraph g {" ] ++ xs ++ [ "}", "" ] | |
| invokeDot :: Graph String -> IO FilePath | |
| invokeDot g = do | |
| fpin <- writeSystemTempFile "graph.dot" (graphviz g) | |
| fpout <- emptySystemTempFile "graph.png" | |
| callCommand $ "dot -Tpng " ++ fpin ++ " -o " ++ fpout | |
| return fpout | |
| viewDot :: FilePath -> IO () | |
| viewDot fp = callCommand $ "xdg-open " ++ fp | |
| -- | Visualise the graph by calling out to the @dot@ command | |
| -- and then using @xdg-open@ to open an image viewer. | |
| viewGraph :: Graph String -> IO () | |
| viewGraph = invokeDot >=> viewDot |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment