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