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

main :: IO ()
main = do
    l <- zip [0..] . scanl (flip Tree.insert) Tree.empty <$> shuffleM [1..150]
    forM_ l $ \(index, tree) -> do
        let (nodes, edges, ranks) = tree2dot 0 Nothing tree
        withFile ("state" ++ show index ++ ".dot") WriteMode $ \file -> do
            hPutStrLn file "digraph tree {"
            hPutStrLn file "nodesep = 0.1;"
            hPutStrLn file "ranksep = 0.5;"
            hPutStrLn file "node [height=0.3, width=0.3, margin=0.05];"
            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" ++) . show) l) ++ ";"
                hPutStrLn file "}"
            forM_ nodes $ \(value, color) -> do
                let attrs =
                        (if color then ["color = \"red\""] else []) ++
                        ["label = \"" ++ show value ++ "\""]
                hPutStrLn file $ "node" ++ show 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" ++ show fnode' ++ " -> node" ++
                    show tnode' ++ "[" ++ intercalate ", " attrs ++ "];"
            hPutStrLn file "}"
        system $ "dot -Tpdf -o state" ++ show index ++
            ".pdf state" ++ show index ++ ".dot"
    system $ "pdftk " ++
        unwords (Prelude.map (("state" ++) . (++ ".pdf") . show) [0..150]) ++
        " cat output all.pdf"
    return ()