Clean up Typeable; derive more Generic
[packages/containers.git] / Data / Tree.hs
1 {-# LANGUAGE CPP #-}
2 #if __GLASGOW_HASKELL__
3 {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-}
4 #endif
5 #if __GLASGOW_HASKELL__ >= 702
6 {-# LANGUAGE DeriveGeneric #-}
7 #endif
8 #if __GLASGOW_HASKELL__ >= 703
9 {-# LANGUAGE Trustworthy #-}
10 #endif
11
12 #include "containers.h"
13
14 -----------------------------------------------------------------------------
15 -- |
16 -- Module : Data.Tree
17 -- Copyright : (c) The University of Glasgow 2002
18 -- License : BSD-style (see the file libraries/base/LICENSE)
19 --
20 -- Maintainer : libraries@haskell.org
21 -- Stability : experimental
22 -- Portability : portable
23 --
24 -- Multi-way trees (/aka/ rose trees) and forests.
25 --
26 -----------------------------------------------------------------------------
27
28 module Data.Tree(
29 Tree(..), Forest,
30 -- * Two-dimensional drawing
31 drawTree, drawForest,
32 -- * Extraction
33 flatten, levels, foldTree,
34 -- * Building trees
35 unfoldTree, unfoldForest,
36 unfoldTreeM, unfoldForestM,
37 unfoldTreeM_BF, unfoldForestM_BF,
38 ) where
39
40 #if MIN_VERSION_base(4,8,0)
41 import Data.Foldable (toList)
42 #else
43 import Control.Applicative (Applicative(..), (<$>))
44 import Data.Foldable (Foldable(foldMap), toList)
45 import Data.Monoid (Monoid(..))
46 import Data.Traversable (Traversable(traverse))
47 #endif
48
49 import Control.Monad (liftM)
50 import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList,
51 ViewL(..), ViewR(..), viewl, viewr)
52 import Data.Typeable
53 import Control.DeepSeq (NFData(rnf))
54
55 #ifdef __GLASGOW_HASKELL__
56 import Data.Data (Data)
57 #endif
58 #if __GLASGOW_HASKELL__ >= 706
59 import GHC.Generics (Generic, Generic1)
60 #elif __GLASGOW_HASKELL__ >= 702
61 import GHC.Generics (Generic)
62 #endif
63
64 #if MIN_VERSION_base(4,8,0)
65 import Data.Coerce
66 #endif
67
68 -- | Multi-way trees, also known as /rose trees/.
69 data Tree a = Node {
70 rootLabel :: a, -- ^ label value
71 subForest :: Forest a -- ^ zero or more child trees
72 }
73 #ifdef __GLASGOW_HASKELL__
74 #if __GLASGOW_HASKELL__ >= 706
75 deriving (Eq, Read, Show, Data, Generic, Generic1)
76 #elif __GLASGOW_HASKELL__ >= 702
77 deriving (Eq, Read, Show, Data, Generic)
78 #else
79 deriving (Eq, Read, Show, Data)
80 #endif
81 #else
82 deriving (Eq, Read, Show)
83 #endif
84 type Forest a = [Tree a]
85
86 INSTANCE_TYPEABLE1(Tree)
87
88 instance Functor Tree where
89 fmap = fmapTree
90
91 fmapTree :: (a -> b) -> Tree a -> Tree b
92 fmapTree f (Node x ts) = Node (f x) (map (fmapTree f) ts)
93 #if MIN_VERSION_base(4,8,0)
94 -- Safe coercions were introduced in 4.7.0, but I am not sure if they played
95 -- well enough with RULES to do what we want.
96 {-# NOINLINE [1] fmapTree #-}
97 {-# RULES
98 "fmapTree/coerce" fmapTree coerce = coerce
99 #-}
100 #endif
101
102 instance Applicative Tree where
103 pure x = Node x []
104 Node f tfs <*> tx@(Node x txs) =
105 Node (f x) (map (f <$>) txs ++ map (<*> tx) tfs)
106
107 instance Monad Tree where
108 return = pure
109 Node x ts >>= f = Node x' (ts' ++ map (>>= f) ts)
110 where Node x' ts' = f x
111
112 instance Traversable Tree where
113 traverse f (Node x ts) = Node <$> f x <*> traverse (traverse f) ts
114
115 instance Foldable Tree where
116 foldMap f (Node x ts) = f x `mappend` foldMap (foldMap f) ts
117
118 #if MIN_VERSION_base(4,8,0)
119 null _ = False
120 {-# INLINE null #-}
121 toList = flatten
122 {-# INLINE toList #-}
123 #endif
124
125 instance NFData a => NFData (Tree a) where
126 rnf (Node x ts) = rnf x `seq` rnf ts
127
128 -- | Neat 2-dimensional drawing of a tree.
129 drawTree :: Tree String -> String
130 drawTree = unlines . draw
131
132 -- | Neat 2-dimensional drawing of a forest.
133 drawForest :: Forest String -> String
134 drawForest = unlines . map drawTree
135
136 draw :: Tree String -> [String]
137 draw (Node x ts0) = lines x ++ drawSubTrees ts0
138 where
139 drawSubTrees [] = []
140 drawSubTrees [t] =
141 "|" : shift "`- " " " (draw t)
142 drawSubTrees (t:ts) =
143 "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts
144
145 shift first other = zipWith (++) (first : repeat other)
146
147 -- | The elements of a tree in pre-order.
148 flatten :: Tree a -> [a]
149 flatten t = squish t []
150 where squish (Node x ts) xs = x:Prelude.foldr squish xs ts
151
152 -- | Lists of nodes at each level of the tree.
153 levels :: Tree a -> [[a]]
154 levels t =
155 map (map rootLabel) $
156 takeWhile (not . null) $
157 iterate (concatMap subForest) [t]
158
159 -- | Catamorphism on trees.
160 foldTree :: (a -> [b] -> b) -> Tree a -> b
161 foldTree f = go where
162 go (Node x ts) = f x (map go ts)
163
164 -- | Build a tree from a seed value
165 unfoldTree :: (b -> (a, [b])) -> b -> Tree a
166 unfoldTree f b = let (a, bs) = f b in Node a (unfoldForest f bs)
167
168 -- | Build a forest from a list of seed values
169 unfoldForest :: (b -> (a, [b])) -> [b] -> Forest a
170 unfoldForest f = map (unfoldTree f)
171
172 -- | Monadic tree builder, in depth-first order
173 unfoldTreeM :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
174 unfoldTreeM f b = do
175 (a, bs) <- f b
176 ts <- unfoldForestM f bs
177 return (Node a ts)
178
179 -- | Monadic forest builder, in depth-first order
180 unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
181 unfoldForestM f = Prelude.mapM (unfoldTreeM f)
182
183 -- | Monadic tree builder, in breadth-first order,
184 -- using an algorithm adapted from
185 -- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
186 -- by Chris Okasaki, /ICFP'00/.
187 unfoldTreeM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
188 unfoldTreeM_BF f b = liftM getElement $ unfoldForestQ f (singleton b)
189 where
190 getElement xs = case viewl xs of
191 x :< _ -> x
192 EmptyL -> error "unfoldTreeM_BF"
193
194 -- | Monadic forest builder, in breadth-first order,
195 -- using an algorithm adapted from
196 -- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
197 -- by Chris Okasaki, /ICFP'00/.
198 unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
199 unfoldForestM_BF f = liftM toList . unfoldForestQ f . fromList
200
201 -- takes a sequence (queue) of seeds
202 -- produces a sequence (reversed queue) of trees of the same length
203 unfoldForestQ :: Monad m => (b -> m (a, [b])) -> Seq b -> m (Seq (Tree a))
204 unfoldForestQ f aQ = case viewl aQ of
205 EmptyL -> return empty
206 a :< aQ' -> do
207 (b, as) <- f a
208 tQ <- unfoldForestQ f (Prelude.foldl (|>) aQ' as)
209 let (tQ', ts) = splitOnto [] as tQ
210 return (Node b ts <| tQ')
211 where
212 splitOnto :: [a'] -> [b'] -> Seq a' -> (Seq a', [a'])
213 splitOnto as [] q = (q, as)
214 splitOnto as (_:bs) q = case viewr q of
215 q' :> a -> splitOnto (a:as) bs q'
216 EmptyR -> error "unfoldForestQ"