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 -----------------------------------------------------------------------------
  14
  15 module Data.Tree(
  16         Tree(..), Forest,
  17         drawTree, drawForest,
  18         flatten, levels,
  19     ) where
  20
  21 #ifdef __HADDOCK__
  22 import Prelude
  23 #endif
  24
  25 -- | Multi-way trees, also known as /rose trees/.
  26 data Tree a   = Node a (Forest a) -- ^ a value and zero or more child trees.
  27 #ifndef __HADDOCK__
  28   deriving (Eq, Read, Show)
  29 #else /* __HADDOCK__ (which can't figure these out by itself) */
  30 instance Eq a => Eq (Tree a)
  31 instance Read a => Read (Tree a)
  32 instance Show a => Show (Tree a)
  33 #endif
  34 type Forest a = [Tree a]
  35
  36 instance Functor Tree where
  37   fmap = mapTree
  38
  39 mapTree              :: (a -> b) -> (Tree a -> Tree b)
  40 mapTree f (Node x ts) = Node (f x) (map (mapTree f) ts)
  41
  42 -- | Neat 2-dimensional drawing of a tree.
  43 drawTree :: Show a => Tree a -> String
  44 drawTree  = unlines . draw . mapTree show
  45
  46 -- | Neat 2-dimensional drawing of a forest.
  47 drawForest :: Show a => Forest a -> String
  48 drawForest  = unlines . map drawTree
  49
  50 draw :: Tree String -> [String]
  51 draw (Node x ts0) = grp this (space (length this)) (stLoop ts0)
  52  where this          = s1 ++ x ++ " "
  53
  54        space n       = replicate n ' '
  55
  56        stLoop []     = [""]
  57        stLoop [t]    = grp s2 "  " (draw t)
  58        stLoop (t:ts) = grp s3 s4 (draw t) ++ [s4] ++ rsLoop ts
  59
  60        rsLoop []     = error "rsLoop:Unexpected empty list."
  61        rsLoop [t]    = grp s5 "  " (draw t)
  62        rsLoop (t:ts) = grp s6 s4 (draw t) ++ [s4] ++ rsLoop ts
  63
  64        grp fst0 rst  = zipWith (++) (fst0:repeat rst)
  65
  66        [s1,s2,s3,s4,s5,s6] = ["- ", "--", "-+", " |", " `", " +"]
  67
  68 -- | The elements of a tree in pre-order.
  69 flatten :: Tree a -> [a]
  70 flatten t = squish t []
  71  where squish (Node x ts) xs = x:foldr squish xs ts
  72
  73 -- | Lists of nodes at each level of the tree.
  74 levels :: Tree a -> [[a]]
  75 levels t = map (map root) $ takeWhile (not . null) $ iterate subforest [t]
  76  where root (Node x _) = x
  77        subforest f     = [t | Node _ ts <- f, t <- ts]