Data/Tree.hs

   1 -----------------------------------------------------------------------------
   2 -- |
   3 -- Module      :  Data.Tree
   4 -- Copyright   :  (c) The University of Glasgow 2002
   5 -- License     :  BSD-style (see the file libraries/base/LICENSE)
   6 --
   7 -- Maintainer  :  libraries@haskell.org
   8 -- Stability   :  experimental
   9 -- Portability :  portable
  10 --
  11 -- Multi-way trees (/aka/ rose trees) and forests.
  12 --
  13 -- Also included are neat presentations for trees and forests.
  14 --
  15 -----------------------------------------------------------------------------
  16
  17 module Data.Tree(
  18         Tree(..), Forest,
  19         flatten, levels,
  20     ) where
  21
  22 -- | Multi-way trees, also known as /rose trees/.
  23 data Tree a   = Node a (Forest a) -- ^ a value and zero or more child trees.
  24 type Forest a = [Tree a]
  25
  26 instance Functor Tree where
  27   fmap = mapTree
  28
  29 mapTree              :: (a -> b) -> (Tree a -> Tree b)
  30 mapTree f (Node x ts) = Node (f x) (map (mapTree f) ts)
  31
  32 instance Show a => Show (Tree a) where
  33   show = showTree
  34   showList ts s = showForest ts ++ s
  35
  36 showTree :: Show a => Tree a -> String
  37 showTree  = drawTree . mapTree show
  38
  39 showForest :: Show a => Forest a -> String
  40 showForest  = unlines . map showTree
  41
  42 drawTree :: Tree String -> String
  43 drawTree  = unlines . draw
  44
  45 draw :: Tree String -> [String]
  46 draw (Node x ts0) = grp this (space (length this)) (stLoop ts0)
  47  where this          = s1 ++ x ++ " "
  48
  49        space n       = replicate n ' '
  50
  51        stLoop []     = [""]
  52        stLoop [t]    = grp s2 "  " (draw t)
  53        stLoop (t:ts) = grp s3 s4 (draw t) ++ [s4] ++ rsLoop ts
  54
  55        rsLoop []     = error "rsLoop:Unexpected empty list."
  56        rsLoop [t]    = grp s5 "  " (draw t)
  57        rsLoop (t:ts) = grp s6 s4 (draw t) ++ [s4] ++ rsLoop ts
  58
  59        grp fst0 rst  = zipWith (++) (fst0:repeat rst)
  60
  61        [s1,s2,s3,s4,s5,s6] = ["- ", "--", "-+", " |", " `", " +"]
  62
  63 -- | The elements of a tree in pre-order.
  64 flatten :: Tree a -> [a]
  65 flatten t = squish t []
  66  where squish (Node x ts) xs = x:foldr squish xs ts
  67
  68 -- | Lists of nodes at each level of the tree.
  69 levels :: Tree a -> [[a]]
  70 levels t = map (map root) $ takeWhile (not . null) $ iterate subforest [t]
  71  where root (Node x _) = x
  72        subforest f     = [t | Node _ ts <- f, t <- ts]