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 -- explicit instance for Haddock's benefit
  33 instance Eq a => Eq (Tree a) where
  34   Node x ts == Node x' ts'  =  x == x' && ts == ts'
  35
  36 instance Show a => Show (Tree a) where
  37   show = showTree
  38   showList ts s = showForest ts ++ s
  39
  40 showTree :: Show a => Tree a -> String
  41 showTree  = drawTree . mapTree show
  42
  43 showForest :: Show a => Forest a -> String
  44 showForest  = unlines . map showTree
  45
  46 drawTree :: Tree String -> String
  47 drawTree  = unlines . draw
  48
  49 draw :: Tree String -> [String]
  50 draw (Node x ts0) = grp this (space (length this)) (stLoop ts0)
  51  where this          = s1 ++ x ++ " "
  52
  53        space n       = replicate n ' '
  54
  55        stLoop []     = [""]
  56        stLoop [t]    = grp s2 "  " (draw t)
  57        stLoop (t:ts) = grp s3 s4 (draw t) ++ [s4] ++ rsLoop ts
  58
  59        rsLoop []     = error "rsLoop:Unexpected empty list."
  60        rsLoop [t]    = grp s5 "  " (draw t)
  61        rsLoop (t:ts) = grp s6 s4 (draw t) ++ [s4] ++ rsLoop ts
  62
  63        grp fst0 rst  = zipWith (++) (fst0:repeat rst)
  64
  65        [s1,s2,s3,s4,s5,s6] = ["- ", "--", "-+", " |", " `", " +"]
  66
  67 -- | The elements of a tree in pre-order.
  68 flatten :: Tree a -> [a]
  69 flatten t = squish t []
  70  where squish (Node x ts) xs = x:foldr squish xs ts
  71
  72 -- | Lists of nodes at each level of the tree.
  73 levels :: Tree a -> [[a]]
  74 levels t = map (map root) $ takeWhile (not . null) $ iterate subforest [t]
  75  where root (Node x _) = x
  76        subforest f     = [t | Node _ ts <- f, t <- ts]