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