[project @ 2003-11-06 12:50:22 by ross]
[packages/containers.git] / 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]