[Haskell-cafe] Draw K-ary forest in dot script

[larry.liuxinyu]

Hi,

I wrote a Haskell program to parse K-ary forest and convert it to dot script 
(Graphviz).

Here is the literate program.

-- First is some stuff imported:
module Main where

import System.Environment (getArgs)
import Text.ParserCombinators.Parsec
import Control.Monad (mapM_)
import Data.List (concatMap, intercalate)
import System.IO (writeFile)
import Data.Char (isSpace)

-- For each tree in the forest, it is described in pre-order.
-- Example description string of a forest of CLRS[1] Figure 19.5(a):
--   (12), (7, (25)), (15, (28, (41)), (33))

-- Definition of K-ary node
data Node a = Node { root :: a 
                   , children :: [Node a]} deriving (Eq, Show)

-- Definition of Forest
type Forest a = [Node a]


-- parsers

-- a forest is a list of trees separate by ','
forest = do 
  ts <- node `sepBy` (char ',')
  return ts

-- a node contains a key then followed by a children forest or nothing (leaf 
case)
node = do
  char '('
  elem <- key
  ts <- (try (char ',')>>forest) <|> return []
  char ')'
  return (Node elem ts)

-- a key is just a plain literate string.
key = many (noneOf ",()")

-- Command line arguments handling
parseArgs :: [String] -> (String, String)
parseArgs [fname, s] = (fname, s)
parseArgs _ = error "wrong usage\nexample:\nfr2dot output.dot \"(12), (7, 
(25)), (15, ((28, (41)), 33))\""


-- A simplified function to generate dot script from parsed result.
toDot f = forestToDot f "t" True

-- a handy function to convert children of a K-ary tree to dot script
treesToDot ts prefix = forestToDot ts prefix False

-- convert a forest to dot script
forestToDot []  _ _ = ""
forestToDot [t] prefix _ = nodeToDot t prefix
forestToDot ts@(_:_:_) prefix lnk = 
    (concatMap (\t->nodeToDot t prefix) ts) ++ consRoot
    where
      consRoot = "{rank=same " ++ ns ++ vis ++ "}\n" 
      ns = intercalate "->" $ map (\t -> prefix ++ root t) ts
      vis = if lnk then "" else "[style=invis]"


-- convert a node to dot script
nodeToDot (Node x ts) prefix = 
    prefix'++"[label=\""++x++"\"];\n" ++
    (treesToDot ts prefix') ++
    (defCons ts prefix')
        where prefix' = prefix ++ x

-- define connections among nodes in dot format
defCons ts prefix = concatMap f ts where
    f (Node x _) = prefix++"->"++prefix++x++";\n"

-- generate dot script from a parsed forest
genDot fname (Right f) = writeFile fname dots >> putStrLn dots
    where
      dots = "digraph G{\n\tnode[shape=circle]\n"++(addTab $ toDot f)++"}"
      addTab s = unlines $ map ("\t"++) (lines s)

main = do
  args <- getArgs
  let (fname, s) = parseArgs args
  genDot fname (parse forest "unknown" (filter (not.isSpace) s))

-- END

I tested with the following simple cases:
./fr2dot foo.dot "(12), (7, (25)), (15, (28, (41)), (33))"
./fr2dot bar.dot "(18), (3, (37)), (6, (8, (30, (45, (55)), (32)), (23, 
(24)), (22)), (29, (48, (50)), (31)), (10, (17)), (44))"

Run the following commands can convert to PNG files:
./dot -Tpng -o foo.png foo.dot
./dot -Tpng -o bar.png bar.dot

Reference:

[1]. Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest and Clifford 
Stein. ``Introduction to Algorithms, Second Edition''. The MIT Press, 2001. 
ISBN: 0262032937.

Best regards.
--
Larry, LIU
https://sites.google.com/site/algoxy/home
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://www.haskell.org/pipermail/haskell-cafe/attachments/20110109/cacce0d8/attachment.htm>