visualizer.hs

module Main where

import Data.List
import Data.IntMap as IntMap
import Data.Set.RBTree as Tree
import Control.Applicative
import Control.Monad
import System.IO
import System.Random.Shuffle
import System.Process

tree2dot :: Int -> Maybe a -> RBTree a ->
    ([(a, Bool)], [(a, a, Bool)], IntMap [a])
tree2dot _ _ Leaf = ([], [], IntMap.empty)
tree2dot n p (Node color _ l v r) =
    (n1 ++ [(v, flag)] ++ n2,
     e1 ++ maybe [] (\p -> [(p, v, flag)]) p ++ e2,
     IntMap.unionWith (++) r1 (IntMap.insertWith (++) n' [v] r2))
    where
    (flag, n') = case color of
        B -> (False, succ n)
        R -> (True, n)
    (n1, e1, r1) = tree2dot n' (Just v) l
    (n2, e2, r2) = tree2dot n' (Just v) r

putDot :: Ord a => (a -> String) -> Handle ->
          ([(a, Bool)], [(a, a, Bool)], IntMap [a]) -> IO ()
putDot toStr file (nodes, edges, ranks) = do
    hPutStrLn file "digraph tree {"
    hPutStrLn file "nodesep = 0.0;"
    hPutStrLn file "ranksep = 0.5;"
    hPutStrLn file "node [height=0.3, width=0.3, margin=0];"
    hPutStrLn file "edge [arrowsize=0.5];"
    forM_ (IntMap.toList ranks) $ \(_, l) -> do
        hPutStrLn file "{"
        hPutStrLn file "rank = same;"
        hPutStrLn file "edge [style = invis];"
        hPutStrLn file $ intercalate " -> "
            (Prelude.map (("node" ++) . toStr) l) ++ ";"
        hPutStrLn file "}"
    forM_ nodes $ \(value, color) -> do
        let attrs =
                (if color then ["color = \"red\""] else []) ++
                ["label = \"" ++ toStr value ++ "\""]
        hPutStrLn file $ "node" ++ toStr value ++
            "[" ++ intercalate ", " attrs ++ "];"
    forM_ edges $ \(fnode, tnode, color) -> do
        let attrs =
                (if color then ["color = \"red\""] else []) ++
                (if tnode < fnode then ["dir = back"] else [])
        let fnode' = min fnode tnode
        let tnode' = max fnode tnode
        hPutStrLn file $ "node" ++ toStr fnode' ++ " -> node" ++
            toStr tnode' ++ "[" ++ intercalate ", " attrs ++ "];"
    hPutStrLn file "}"

main :: IO ()
main = do
    l <- zip [0..] . scanl (flip Tree.insert) Tree.empty <$> shuffleM [1..100]
    forM_ l $ \(index, tree) -> do
        withFile ("state" ++ show index ++ ".dot") WriteMode $
            flip (putDot show) (tree2dot 0 Nothing tree)
        system $ "dot -Tpdf -o state" ++ show index ++
            ".pdf state" ++ show index ++ ".dot"
    system $ "pdftk " ++
        unwords (Prelude.map (("state" ++) . (++ ".pdf") . show) [0..100]) ++
        " cat output all.pdf"
    return ()