3f3a1bb5a1ec8e31e3dd4aef23a756d8b9e9cd55
[packages/containers.git] / Data / Tree.hs
1 {-# LANGUAGE PatternGuards #-}
2 {-# LANGUAGE CPP #-}
3 #if __GLASGOW_HASKELL__
4 {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-}
5 {-# LANGUAGE DeriveGeneric #-}
6 {-# LANGUAGE Trustworthy #-}
7 #endif
8
9 #include "containers.h"
10
11 -----------------------------------------------------------------------------
12 -- |
13 -- Module : Data.Tree
14 -- Copyright : (c) The University of Glasgow 2002
15 -- License : BSD-style (see the file libraries/base/LICENSE)
16 --
17 -- Maintainer : libraries@haskell.org
18 -- Portability : portable
19 --
20 -- = Multi-way Trees and Forests
21 --
22 -- The @'Tree' a@ type represents a lazy, possibly infinite, multi-way tree
23 -- (also known as a /rose tree/).
24 --
25 -- The @'Forest' a@ type represents a forest of @'Tree' a@s.
26 --
27 -----------------------------------------------------------------------------
28
29 module Data.Tree(
30
31 -- * Trees and Forests
32 Tree(..)
33 , Forest
34
35 -- * Construction
36 , unfoldTree
37 , unfoldForest
38 , unfoldTreeM
39 , unfoldForestM
40 , unfoldTreeM_BF
41 , unfoldForestM_BF
42
43 -- * Elimination
44 , foldTree
45 , flatten
46 , levels
47
48 -- * Ascii Drawings
49 , drawTree
50 , drawForest
51
52 ) where
53
54 #if MIN_VERSION_base(4,8,0)
55 import Data.Foldable (toList)
56 import Control.Applicative (Applicative(..), liftA2)
57 #else
58 import Control.Applicative (Applicative(..), liftA2, (<$>))
59 import Data.Foldable (Foldable(foldMap), toList)
60 import Data.Monoid (Monoid(..))
61 import Data.Traversable (Traversable(traverse))
62 #endif
63
64 import Control.Monad (liftM)
65 import Control.Monad.Fix (MonadFix (..), fix)
66 import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList,
67 ViewL(..), ViewR(..), viewl, viewr)
68 import Data.Typeable
69 import Control.DeepSeq (NFData(rnf))
70
71 #ifdef __GLASGOW_HASKELL__
72 import Data.Data (Data)
73 import GHC.Generics (Generic, Generic1)
74 #endif
75
76 import Control.Monad.Zip (MonadZip (..))
77
78 #if MIN_VERSION_base(4,8,0)
79 import Data.Coerce
80 #endif
81
82 #if MIN_VERSION_base(4,9,0)
83 import Data.Functor.Classes
84 import Data.Semigroup (Semigroup (..))
85 #endif
86
87 #if !MIN_VERSION_base(4,8,0)
88 import Data.Functor ((<$))
89 #endif
90
91 -- | Non-empty, possibly infinite, multi-way trees; also known as /rose trees/.
92 data Tree a = Node {
93 rootLabel :: a, -- ^ label value
94 subForest :: Forest a -- ^ zero or more child trees
95 }
96 #ifdef __GLASGOW_HASKELL__
97 deriving ( Eq
98 , Read
99 , Show
100 , Data
101 , Generic -- ^ @since 0.5.8
102 , Generic1 -- ^ @since 0.5.8
103 )
104 #else
105 deriving (Eq, Read, Show)
106 #endif
107
108 type Forest a = [Tree a]
109
110 #if MIN_VERSION_base(4,9,0)
111 -- | @since 0.5.9
112 instance Eq1 Tree where
113 liftEq eq = leq
114 where
115 leq (Node a fr) (Node a' fr') = eq a a' && liftEq leq fr fr'
116
117 -- | @since 0.5.9
118 instance Ord1 Tree where
119 liftCompare cmp = lcomp
120 where
121 lcomp (Node a fr) (Node a' fr') = cmp a a' <> liftCompare lcomp fr fr'
122
123 -- | @since 0.5.9
124 instance Show1 Tree where
125 liftShowsPrec shw shwl p (Node a fr) = showParen (p > 10) $
126 showString "Node {rootLabel = " . shw 0 a . showString ", " .
127 showString "subForest = " . liftShowList shw shwl fr .
128 showString "}"
129
130 -- | @since 0.5.9
131 instance Read1 Tree where
132 liftReadsPrec rd rdl p = readParen (p > 10) $
133 \s -> do
134 ("Node", s1) <- lex s
135 ("{", s2) <- lex s1
136 ("rootLabel", s3) <- lex s2
137 ("=", s4) <- lex s3
138 (a, s5) <- rd 0 s4
139 (",", s6) <- lex s5
140 ("subForest", s7) <- lex s6
141 ("=", s8) <- lex s7
142 (fr, s9) <- liftReadList rd rdl s8
143 ("}", s10) <- lex s9
144 pure (Node a fr, s10)
145 #endif
146
147 INSTANCE_TYPEABLE1(Tree)
148
149 instance Functor Tree where
150 fmap = fmapTree
151 x <$ Node _ ts = Node x (map (x <$) ts)
152
153 fmapTree :: (a -> b) -> Tree a -> Tree b
154 fmapTree f (Node x ts) = Node (f x) (map (fmapTree f) ts)
155 #if MIN_VERSION_base(4,8,0)
156 -- Safe coercions were introduced in 4.7.0, but I am not sure if they played
157 -- well enough with RULES to do what we want.
158 {-# NOINLINE [1] fmapTree #-}
159 {-# RULES
160 "fmapTree/coerce" fmapTree coerce = coerce
161 #-}
162 #endif
163
164 instance Applicative Tree where
165 pure x = Node x []
166 Node f tfs <*> tx@(Node x txs) =
167 Node (f x) (map (f <$>) txs ++ map (<*> tx) tfs)
168 #if MIN_VERSION_base(4,10,0)
169 liftA2 f (Node x txs) ty@(Node y tys) =
170 Node (f x y) (map (f x <$>) tys ++ map (\tx -> liftA2 f tx ty) txs)
171 #endif
172 Node x txs <* ty@(Node _ tys) =
173 Node x (map (x <$) tys ++ map (<* ty) txs)
174 Node _ txs *> ty@(Node y tys) =
175 Node y (tys ++ map (*> ty) txs)
176
177 instance Monad Tree where
178 return = pure
179 Node x ts >>= f = case f x of
180 Node x' ts' -> Node x' (ts' ++ map (>>= f) ts)
181
182 -- | @since 0.5.11
183 instance MonadFix Tree where
184 mfix = mfixTree
185
186 mfixTree :: (a -> Tree a) -> Tree a
187 mfixTree f
188 | Node a children <- fix (f . rootLabel)
189 = Node a (zipWith (\i _ -> mfixTree ((!! i) . subForest . f))
190 [0..] children)
191
192 instance Traversable Tree where
193 traverse f (Node x ts) = liftA2 Node (f x) (traverse (traverse f) ts)
194
195 instance Foldable Tree where
196 foldMap f (Node x ts) = f x `mappend` foldMap (foldMap f) ts
197
198 #if MIN_VERSION_base(4,8,0)
199 null _ = False
200 {-# INLINE null #-}
201 toList = flatten
202 {-# INLINE toList #-}
203 #endif
204
205 instance NFData a => NFData (Tree a) where
206 rnf (Node x ts) = rnf x `seq` rnf ts
207
208 instance MonadZip Tree where
209 mzipWith f (Node a as) (Node b bs)
210 = Node (f a b) (mzipWith (mzipWith f) as bs)
211
212 munzip (Node (a, b) ts) = (Node a as, Node b bs)
213 where (as, bs) = munzip (map munzip ts)
214
215 -- | 2-dimensional ASCII drawing of a tree.
216 --
217 -- ==== __Examples__
218 --
219 -- > putStr $ drawTree $ fmap show (Node 1 [Node 2 [], Node 3 []])
220 --
221 -- @
222 -- 1
223 -- |
224 -- +- 2
225 -- |
226 -- `- 3
227 -- @
228 --
229 drawTree :: Tree String -> String
230 drawTree = unlines . draw
231
232 -- | 2-dimensional ASCII drawing of a forest.
233 --
234 -- ==== __Examples__
235 --
236 -- > putStr $ drawForest $ map (fmap show) [(Node 1 [Node 2 [], Node 3 []]), (Node 10 [Node 20 []])]
237 --
238 -- @
239 -- 1
240 -- |
241 -- +- 2
242 -- |
243 -- `- 3
244 --
245 -- 10
246 -- |
247 -- `- 20
248 -- @
249 --
250 drawForest :: Forest String -> String
251 drawForest = unlines . map drawTree
252
253 draw :: Tree String -> [String]
254 draw (Node x ts0) = lines x ++ drawSubTrees ts0
255 where
256 drawSubTrees [] = []
257 drawSubTrees [t] =
258 "|" : shift "`- " " " (draw t)
259 drawSubTrees (t:ts) =
260 "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts
261
262 shift first other = zipWith (++) (first : repeat other)
263
264 -- | Returns the elements of a tree in pre-order.
265 --
266 -- @
267 --
268 -- a
269 -- / \\ => [a,b,c]
270 -- b c
271 -- @
272 --
273 -- ==== __Examples__
274 --
275 -- > flatten (Node 1 [Node 2 [], Node 3 []]) == [1,2,3]
276 flatten :: Tree a -> [a]
277 flatten t = squish t []
278 where squish (Node x ts) xs = x:Prelude.foldr squish xs ts
279
280 -- | Returns the list of nodes at each level of the tree.
281 --
282 -- @
283 --
284 -- a
285 -- / \\ => [[a], [b,c]]
286 -- b c
287 -- @
288 --
289 -- ==== __Examples__
290 --
291 -- > levels (Node 1 [Node 2 [], Node 3 []]) == [[1],[2,3]]
292 --
293 levels :: Tree a -> [[a]]
294 levels t =
295 map (map rootLabel) $
296 takeWhile (not . null) $
297 iterate (concatMap subForest) [t]
298
299 -- | Fold a tree into a "summary" value in depth-first order.
300 --
301 -- For each node in the tree, apply @f@ to the @rootLabel@ and the result
302 -- of applying @f@ to each @subForent@.
303 --
304 -- This is also known as the catamorphism on trees.
305 --
306 -- ==== __Examples__
307 --
308 -- Sum the values in a tree:
309 --
310 -- > foldTree (\x xs -> sum (x:xs)) (Node 1 [Node 2 [], Node 3 []]) == 6
311 --
312 -- Find the maximum value in the tree:
313 --
314 -- > foldTree (\x xs -> maximum (x:xs)) (Node 1 [Node 2 [], Node 3 []]) == 3
315 --
316 --
317 -- @since 0.5.8
318 foldTree :: (a -> [b] -> b) -> Tree a -> b
319 foldTree f = go where
320 go (Node x ts) = f x (map go ts)
321
322 -- | Build a (possibly infinite) tree from a seed value in breadth-first order.
323 --
324 -- @unfoldTree f b@ constructs a tree by starting with the tree
325 -- @Node { rootLabel=b, subForest=[] }@ and repeatedly applying @f@ to each
326 -- 'rootLabel' value in the tree's leaves to generate its 'subForest'.
327 --
328 -- For a monadic version see 'unfoldTreeM_BF'.
329 --
330 -- ==== __Examples__
331 --
332 -- Construct the tree of @Integer@s where each node has two children:
333 -- @left = 2*x@ and @right = 2*x + 1@, where @x@ is the 'rootLabel' of the node.
334 -- Stop when the values exceed 7.
335 --
336 -- > let buildNode x = if 2*x + 1 > 7 then (x, []) else (x, [2*x, 2*x+1])
337 -- > putStr $ drawTree $ fmap show $ unfoldTree buildNode 1
338 --
339 -- @
340 --
341 -- 1
342 -- |
343 -- +- 2
344 -- | |
345 -- | +- 4
346 -- | |
347 -- | `- 5
348 -- |
349 -- `- 3
350 -- |
351 -- +- 6
352 -- |
353 -- `- 7
354 -- @
355 --
356 unfoldTree :: (b -> (a, [b])) -> b -> Tree a
357 unfoldTree f b = let (a, bs) = f b in Node a (unfoldForest f bs)
358
359 -- | Build a (possibly infinite) forest from a list of seed values in
360 -- breadth-first order.
361 --
362 -- @unfoldForest f seeds@ invokes 'unfoldTree' on each seed value.
363 --
364 -- For a monadic version see 'unfoldForestM_BF'.
365 --
366 unfoldForest :: (b -> (a, [b])) -> [b] -> Forest a
367 unfoldForest f = map (unfoldTree f)
368
369 -- | Monadic tree builder, in depth-first order.
370 unfoldTreeM :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
371 unfoldTreeM f b = do
372 (a, bs) <- f b
373 ts <- unfoldForestM f bs
374 return (Node a ts)
375
376 -- | Monadic forest builder, in depth-first order
377 unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
378 unfoldForestM f = Prelude.mapM (unfoldTreeM f)
379
380 -- | Monadic tree builder, in breadth-first order.
381 --
382 -- See 'unfoldTree' for more info.
383 --
384 -- Implemented using an algorithm adapted from /Breadth-First Numbering: Lessons
385 -- from a Small Exercise in Algorithm Design/, by Chris Okasaki, /ICFP'00/.
386 unfoldTreeM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
387 unfoldTreeM_BF f b = liftM getElement $ unfoldForestQ f (singleton b)
388 where
389 getElement xs = case viewl xs of
390 x :< _ -> x
391 EmptyL -> error "unfoldTreeM_BF"
392
393 -- | Monadic forest builder, in breadth-first order
394 --
395 -- See 'unfoldForest' for more info.
396 --
397 -- Implemented using an algorithm adapted from /Breadth-First Numbering: Lessons
398 -- from a Small Exercise in Algorithm Design/, by Chris Okasaki, /ICFP'00/.
399 unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
400 unfoldForestM_BF f = liftM toList . unfoldForestQ f . fromList
401
402 -- Takes a sequence (queue) of seeds and produces a sequence (reversed queue) of
403 -- trees of the same length.
404 unfoldForestQ :: Monad m => (b -> m (a, [b])) -> Seq b -> m (Seq (Tree a))
405 unfoldForestQ f aQ = case viewl aQ of
406 EmptyL -> return empty
407 a :< aQ' -> do
408 (b, as) <- f a
409 tQ <- unfoldForestQ f (Prelude.foldl (|>) aQ' as)
410 let (tQ', ts) = splitOnto [] as tQ
411 return (Node b ts <| tQ')
412 where
413 splitOnto :: [a'] -> [b'] -> Seq a' -> (Seq a', [a'])
414 splitOnto as [] q = (q, as)
415 splitOnto as (_:bs) q = case viewr q of
416 q' :> a -> splitOnto (a:as) bs q'
417 EmptyR -> error "unfoldForestQ"