Draw.hs

-- | Generic ascii art \/ graphviz drawing of trees.
--
-- Suggestions for drawing styles are welcome. 
--
-- TODO:
--
--  * make the style customizable
--
--  * the same for graphviz
--
module Data.Generics.Fixplate.Draw
  (
    -- * Default tree drawing, using Show instancess
    drawTree 
  , showTree
  , graphvizTree
    -- * Customizable tree drawing
  , drawTreeWith
  , showTreeWith
  , graphvizTreeWith
  )
  where

--------------------------------------------------------------------------------

import Data.Foldable
import Data.Traversable

import Data.Generics.Fixplate.Base
import Data.Generics.Fixplate.Open

import Data.Generics.Fixplate.Attributes ( enumerateNodes_ )
import Data.Generics.Fixplate.Traversals ( universe )

--------------------------------------------------------------------------------

{-
-- | This a data type defined to be a place-holder for childs.
-- So that you can define it to be an instance of your own pretty-printer.
--
-- For the fastest result, you want to define something like
-- 
-- > instance Show Hole where show _ = "_"
--
-- We don't do this so that you can customize to your preferred drawing style.
-- However, `drawTree' and `showTree' does exactly this.
data Hole = Hole
-}

--------------------------------------------------------------------------------

-- | Prints a tree. It is defined simply as
--
-- > drawTree = putStrLn . showTree
--
drawTree :: (Functor f, Foldable f, ShowF f) => Mu f -> IO ()
drawTree = putStrLn . showTree

drawTreeWith :: (Functor f, Foldable f) => (f Hole -> String) -> Mu f -> IO ()
drawTreeWith pp = putStrLn . showTreeWith pp

--------------------------------------------------------------------------------

type Step = [Bool]

--------------------------------------------------------------------------------

-- this is distinct from Hole so that we that user can defined his own 'Show' instnace for 'Hole'
data Void = Void ; instance Show Void where show _ = "_"

-- | Creates a string representation which can be printed with 'putStrLn'.
showTree :: (Functor f, Foldable f, ShowF f) => Mu f -> String
showTree = showTreeWith pp where
  pp t = showF (fmap (const Void) t) 

--------------------------------------------------------------------------------

showTreeWith :: (Functor f, Foldable f) => (f Hole -> String) -> Mu f -> String
showTreeWith pprint = unlines . map mkLine . go [False] where

--  go :: Step -> Mu f -> [(Step,String)]
  go bars (Fix s) = ( bars , this ) : rest where    
    this = pprint $ fmap (const Hole) s
    rest = Prelude.concat $ reverse $ zipWith worker theBars (toRevList s)
    worker b t = go (b:bars) t 

  theBars = False : repeat True     -- last child is drawn differently when it has subchilds

  mkLine (b:bs, str) = Prelude.concatMap (_branch style) (reverse bs) ++ (_twig style b) ++ str 

  style = defaultStyle 

--------------------------------------------------------------------------------

-- customizable ascii art style

defaultStyle :: Style
defaultStyle = Style 
  { _twigNorm   = " |-- "
  , _twigLast   = " \\-- "
  , _branchNorm = " |   "
  , _branchLast = "     "
  }

{-
someStyle :: Style
someStyle = Style 
  { _twigNorm   = ">- "
  , _twigLast   = "|- "
  , _branchNorm = "|  "
  , _branchLast = "   "
  }
-}

data Style = Style 
  { _twigNorm   :: !String
  , _twigLast   :: !String
  , _branchNorm :: !String
  , _branchLast :: !String
  }

_twig :: Style -> Bool -> String
_twig style b = if b then _twigNorm style  else _twigLast style 

_branch :: Style -> Bool -> String
_branch style b = if b then _branchNorm style  else _branchLast style 


--------------------------------------------------------------------------------

-- | Generate a graphviz @.dot@ file 
graphvizTree :: (Traversable f, ShowF f) => Mu f -> String
graphvizTree = graphvizTreeWith pp where
  pp t = showF (fmap (const Void) t)

graphvizTreeWith :: (Traversable f) => (f Hole -> String) -> Mu f -> String
graphvizTreeWith pp tree = unlines dot where
  dot = header : viznodes ++ vizedges ++ [footer] 
  header = "digraph tree {"
  footer = "}"
  enum = enumerateNodes_ tree
  node i = "node" ++ show i
  only = fmap (const Hole)
  viznodes = 
    [ node i ++ " [ label=\"" ++ escape (pp (only s)) ++ "\" ] ;"
    | Fix (Ann i s) <- universe enum
    ]
  vizedges = Prelude.concat 
    [ [ node i ++ " -> " ++ node j ++ " ;"
      | Fix (Ann j _) <- toList s
      ]
    | Fix (Ann i s) <- universe enum
    ]

escape :: String -> String
escape = Prelude.concatMap f where
  f c = if Prelude.elem c stuff then '\\':c:[] else c:[]
  stuff = "\\\""

--------------------------------------------------------------------------------