{-|
== Draw trees

Primarily intended for debugging.

-}

module TreeDraw
    ( draw
    ) where

-- | Draws a tree. Like Data.Tree.drawTree but works with any type, can show edge labels and uses
-- unicode box characters.
-- Example:
-- a
-- ├── 1 ── b
-- ├── 2 ── c
-- │        ├── d
-- │        └── e
-- └── 3 ── f
draw :: (a -> String) -> (a -> [(Maybe String, a)]) -> a -> String
draw :: (a -> String) -> (a -> [(Maybe String, a)]) -> a -> String
draw a -> String
showa a -> [(Maybe String, a)]
edges = [String] -> String
unlines ([String] -> String) -> (a -> [String]) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [String]
go where
    go :: a -> [String]
go a
a = a -> String
showa a
a String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [(Maybe String, a)] -> [String]
drawCh (a -> [(Maybe String, a)]
edges a
a) where
        drawCh :: [(Maybe String, a)] -> [String]
drawCh []     = []
        drawCh [(Maybe String, a)
e]    = String -> String -> [String] -> [String]
forall a. [a] -> [a] -> [[a]] -> [[a]]
shift String
"└" String
" " ((Maybe String, a) -> [String]
f (Maybe String, a)
e)
        drawCh ((Maybe String, a)
e:[(Maybe String, a)]
es) = String -> String -> [String] -> [String]
forall a. [a] -> [a] -> [[a]] -> [[a]]
shift String
"├" String
"│" ((Maybe String, a) -> [String]
f (Maybe String, a)
e) [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [(Maybe String, a)] -> [String]
drawCh [(Maybe String, a)]
es
        f :: (Maybe String, a) -> [String]
f (Maybe String
e, a
v) = String -> String -> [String] -> [String]
forall a. [a] -> [a] -> [[a]] -> [[a]]
shift String
label String
padding (a -> [String]
go a
v) where
            label :: String
label = Maybe String -> String
edge Maybe String
e
            padding :: String
padding = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
label) Char
' '
    edge :: Maybe String -> String
edge Maybe String
Nothing  = String
"── "
    edge (Just String
s) = String
"── " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" ── "
    shift :: [a] -> [a] -> [[a]] -> [[a]]
shift [a]
first [a]
other = ([a] -> [a] -> [a]) -> [[a]] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) ([a]
first [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
forall a. a -> [a]
repeat [a]
other)