]> Git — Sourcephile - haskell/treemap.git/blob - Data/TreeMap/Strict.hs
Eta-reduce Path type synonym.
[haskell/treemap.git] / Data / TreeMap / Strict.hs
1 {-# LANGUAGE DeriveDataTypeable #-}
2 {-# LANGUAGE NamedFieldPuns #-}
3 {-# LANGUAGE StandaloneDeriving #-}
4 {-# OPTIONS_GHC -fno-warn-orphans #-}
5
6 -- | This module implements a strict 'TreeMap',
7 -- which is like a 'Map'
8 -- but whose key is now a 'NonEmpty' list of 'Map' keys (a 'Path')
9 -- enabling the possibility to gather mapped values
10 -- by 'Path' prefixes (inside a 'Node').
11 module Data.TreeMap.Strict where
12
13 import Control.Applicative (Applicative(..))
14 import Control.DeepSeq (NFData(..))
15 import Data.Bool
16 import Data.Data (Data)
17 import Data.Eq (Eq)
18 import Data.Foldable (Foldable, foldMap)
19 import Data.Function (($), (.), const, flip, id)
20 import Data.Functor (Functor(..), (<$>))
21 import qualified Data.List
22 import qualified Data.List.NonEmpty
23 import Data.List.NonEmpty (NonEmpty(..))
24 import Data.Map.Strict (Map)
25 import qualified Data.Map.Strict as Map
26 import Data.Maybe (Maybe(..), maybe)
27 import Data.Monoid (Monoid(..))
28 import Data.Ord (Ord(..))
29 import qualified Data.Strict.Maybe as Strict
30 import Data.Traversable (Traversable(..))
31 import Data.Typeable (Typeable)
32 import Prelude (Int, Num(..), seq)
33 import Text.Show (Show(..))
34
35 -- @Data.Strict@ orphan instances
36 deriving instance Data x => Data (Strict.Maybe x)
37 deriving instance Typeable Strict.Maybe
38 instance Monoid x => Monoid (Strict.Maybe x) where
39 mempty = Strict.Nothing
40 mappend (Strict.Just x) (Strict.Just y) = Strict.Just (x `mappend` y)
41 mappend x Strict.Nothing = x
42 mappend Strict.Nothing y = y
43 instance NFData x => NFData (Strict.Maybe x) where
44 rnf Strict.Nothing = ()
45 rnf (Strict.Just x) = rnf x
46
47 -- * Type 'TreeMap'
48
49 newtype TreeMap k x
50 = TreeMap (Map k (Node k x))
51 deriving (Data, Eq, Show, Typeable)
52
53 instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
54 mempty = empty
55 mappend = union mappend
56 -- mconcat = Data.List.foldr mappend mempty
57 instance Ord k => Functor (TreeMap k) where
58 fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
59 instance Ord k => Foldable (TreeMap k) where
60 foldMap f (TreeMap m) = foldMap (foldMap f) m
61 instance Ord k => Traversable (TreeMap k) where
62 traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m
63 instance (Ord k, NFData k, NFData x) => NFData (TreeMap k x) where
64 rnf (TreeMap m) = rnf m
65
66 -- * Type 'Path'
67
68 -- | A 'Path' is a non-empty list of 'Map' keys.
69 type Path = NonEmpty
70
71 path :: k -> [k] -> Path k
72 path = (:|)
73
74 list :: Path k -> [k]
75 list = Data.List.NonEmpty.toList
76
77 reverse :: Path k -> Path k
78 reverse = Data.List.NonEmpty.reverse
79
80 -- * Type 'Node'
81 data Node k x
82 = Node
83 { node_size :: !Int -- ^ The number of non-'Strict.Nothing' 'node_value's reachable from this 'Node'.
84 , node_value :: !(Strict.Maybe x) -- ^ Some value, or 'Strict.Nothing' if this 'Node' is intermediary.
85 , node_descendants :: !(TreeMap k x) -- ^ Descendants 'Node's.
86 } deriving (Data, Eq, Show, Typeable)
87
88
89 instance (Ord k, Monoid v) => Monoid (Node k v) where
90 mempty =
91 Node
92 { node_value = Strict.Nothing
93 , node_size = 0
94 , node_descendants = TreeMap mempty
95 }
96 mappend
97 Node{node_value=x0, node_descendants=m0}
98 Node{node_value=x1, node_descendants=m1} =
99 let node_descendants = union const m0 m1 in
100 let node_value = x0 `mappend` x1 in
101 Node
102 { node_value
103 , node_size = size node_descendants
104 + Strict.maybe 0 (const 1) node_value
105 , node_descendants
106 }
107 -- mconcat = Data.List.foldr mappend mempty
108 instance Ord k => Functor (Node k) where
109 fmap f Node{node_value=x, node_descendants=m, node_size} =
110 Node
111 { node_value = fmap f x
112 , node_descendants = map f m
113 , node_size
114 }
115 instance Ord k => Foldable (Node k) where
116 foldMap f Node{node_value=Strict.Nothing, node_descendants=TreeMap m} =
117 foldMap (foldMap f) m
118 foldMap f Node{node_value=Strict.Just x, node_descendants=TreeMap m} =
119 f x `mappend` foldMap (foldMap f) m
120 instance Ord k => Traversable (Node k) where
121 traverse f Node{node_value=Strict.Nothing, node_descendants=TreeMap m, node_size} =
122 Node node_size <$> pure Strict.Nothing <*> (TreeMap <$> traverse (traverse f) m)
123 traverse f Node{node_value=Strict.Just x, node_descendants=TreeMap m, node_size} =
124 Node node_size <$> (Strict.Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
125 instance (Ord k, NFData k, NFData x) => NFData (Node k x) where
126 rnf (Node s v d) = rnf s `seq` rnf v `seq` rnf d
127
128 node_find :: Ord k => [k] -> Node k x -> Strict.Maybe (Node k x)
129 node_find [] n = Strict.Just n
130 node_find (k:ks) Node{node_descendants=TreeMap m} =
131 maybe Strict.Nothing (node_find ks) $
132 Map.lookup k m
133
134 -- * Construct
135
136 -- | Return the empty 'TreeMap'.
137 empty :: TreeMap k x
138 empty = TreeMap Map.empty
139
140 -- | Return a 'TreeMap' only mapping the given 'Path' to the given value.
141 singleton :: Ord k => Path k -> x -> TreeMap k x
142 singleton ks x = insert const ks x empty
143
144 -- | Return a 'Node' only containing the given value.
145 leaf :: x -> Node k x
146 leaf x =
147 Node
148 { node_value = Strict.Just x
149 , node_descendants = empty
150 , node_size = 1
151 }
152
153 -- | Return the given 'TreeMap' associating the given 'Path' with the given value,
154 -- merging values if the given 'TreeMap' already associates the given 'Path'
155 -- with a non-'Strict.Nothing' 'node_value'.
156 insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
157 insert merge (k:|[]) x (TreeMap m) =
158 TreeMap $
159 Map.insertWith
160 (\_ Node{node_value = x1, node_descendants = m1, node_size = s1} ->
161 Node
162 { node_value = Strict.maybe (Strict.Just x) (Strict.Just . merge x) x1
163 , node_descendants = m1
164 , node_size = Strict.maybe (s1 + 1) (const s1) x1
165 })
166 k (leaf x) m
167 insert merge (k:|k':ks) x (TreeMap m) =
168 TreeMap $
169 Map.insertWith
170 (\_ Node{node_value = x1, node_descendants = m1} ->
171 let m' = insert merge (path k' ks) x m1 in
172 let s' = size m' + Strict.maybe 0 (const 1) x1 in
173 Node{node_value=x1, node_descendants=m', node_size=s'})
174 k
175 Node
176 { node_value = Strict.Nothing
177 , node_descendants = insert merge (path k' ks) x empty
178 , node_size = 1
179 }
180 m
181
182 -- | Return a 'TreeMap' associating for each tuple of the given list
183 -- the 'Path' to the value,
184 -- merging values of identical 'Path's (in respective order).
185 from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
186 from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
187
188 -- | Return a 'TreeMap' associating for each key and value of the given 'Map'
189 -- the 'Path' to the value,
190 -- merging values of identical 'Path's (in respective order).
191 from_Map :: Ord k => (x -> x -> x) -> Map (Path k) x -> TreeMap k x
192 from_Map merge = Map.foldlWithKey (\acc p x -> insert merge p x acc) empty
193
194 -- * Size
195
196 -- | Return the 'Map' in the given 'TreeMap'.
197 nodes :: TreeMap k x -> Map k (Node k x)
198 nodes (TreeMap m) = m
199
200 -- | Return 'True' iif. the given 'TreeMap' is 'empty'.
201 null :: TreeMap k x -> Bool
202 null (TreeMap m) = Map.null m
203
204 -- | Return the number of non-'Strict.Nothing' 'node_value's in the given 'TreeMap'.
205 --
206 -- * Complexity: O(r) where r is the size of the root 'Map'.
207 size :: TreeMap k x -> Int
208 size = Map.foldr ((+) . node_size) 0 . nodes
209
210 -- * Find
211
212 -- | Return the value (if any) associated with the given 'Path'.
213 find :: Ord k => Path k -> TreeMap k x -> Strict.Maybe x
214 find (k:|[]) (TreeMap m) = maybe Strict.Nothing node_value $ Map.lookup k m
215 find (k:|k':ks) (TreeMap m) =
216 maybe Strict.Nothing (find (path k' ks) . node_descendants) $
217 Map.lookup k m
218
219 -- | Return the values (if any) associated with the prefixes of the given 'Path' (included).
220 find_along :: Ord k => Path k -> TreeMap k x -> [x]
221 find_along p (TreeMap tm) =
222 go (list p) tm
223 where
224 go :: Ord k => [k] -> Map k (Node k x) -> [x]
225 go [] _m = []
226 go (k:ks) m =
227 case Map.lookup k m of
228 Nothing -> []
229 Just node ->
230 Strict.maybe id (:) (node_value node) $
231 go ks $ nodes (node_descendants node)
232
233 find_node :: Ord k => Path k -> TreeMap k x -> Strict.Maybe (Node k x)
234 find_node (k:|[]) (TreeMap m) = maybe Strict.Nothing Strict.Just $ Map.lookup k m
235 find_node (k:|k':ks) (TreeMap m) =
236 maybe Strict.Nothing (find_node (path k' ks) . node_descendants) $
237 Map.lookup k m
238
239 -- * Union
240
241 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
242 -- merging values (in respective order) when a 'Path' leads
243 -- to a non-'Strict.Nothing' 'node_value' in both given 'TreeMap's.
244 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
245 union merge (TreeMap tm0) (TreeMap tm1) =
246 TreeMap $
247 Map.unionWith
248 (\Node{node_value=x0, node_descendants=m0}
249 Node{node_value=x1, node_descendants=m1} ->
250 let node_descendants = union merge m0 m1 in
251 let node_value = Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0 in
252 Node
253 { node_size = size node_descendants + Strict.maybe 0 (const 1) node_value
254 , node_value
255 , node_descendants
256 })
257 tm0 tm1
258
259 -- | Return the 'union' of the given 'TreeMap's.
260 --
261 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
262 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
263 unions merge = Data.List.foldl' (union merge) empty
264
265 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
266 -- foldl' f = go
267 -- where
268 -- go z [] = z
269 -- go z (x:xs) = z `seq` go (f z x) xs
270
271 -- * Map
272
273 -- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
274 -- mapped by the given function.
275 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
276 map f =
277 TreeMap .
278 Map.map
279 (\n@Node{node_value=x, node_descendants=m} ->
280 n{ node_value = fmap f x
281 , node_descendants = map f m
282 }) .
283 nodes
284
285 -- | Return the given 'TreeMap' with each 'Path' section
286 -- and each non-'Strict.Nothing' 'node_value'
287 -- mapped by the given functions.
288 --
289 -- WARNING: the function mapping 'Path' sections must be monotonic,
290 -- like in 'Map.mapKeysMonotonic'.
291 map_monotonic :: (Ord k, Ord l) => (k -> l) -> (x -> y) -> TreeMap k x -> TreeMap l y
292 map_monotonic fk fx =
293 TreeMap .
294 Map.mapKeysMonotonic fk .
295 Map.map
296 (\n@Node{node_value=x, node_descendants=m} ->
297 n{ node_value = fmap fx x
298 , node_descendants = map_monotonic fk fx m
299 }) .
300 nodes
301
302 -- | Return the given 'TreeMap' with each 'node_value'
303 -- mapped by the given function supplied with
304 -- the already mapped 'node_descendants' of the current 'Node'.
305 map_by_depth_first :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
306 map_by_depth_first f =
307 TreeMap .
308 Map.map
309 (\Node{node_value, node_descendants} ->
310 let m = map_by_depth_first f node_descendants in
311 Node
312 { node_value = Strict.Just $ f m node_value
313 , node_descendants = m
314 , node_size = size m + 1
315 }) .
316 nodes
317
318 -- * Alter
319
320 alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
321 alterl_path fct =
322 go fct . list
323 where
324 go :: Ord k
325 => (Strict.Maybe x -> Strict.Maybe x) -> [k]
326 -> TreeMap k x -> TreeMap k x
327 go _f [] m = m
328 go f (k:p) (TreeMap m) =
329 TreeMap $
330 Map.alter
331 (\c ->
332 let (cv, cm) =
333 case c of
334 Just Node{node_value=v, node_descendants=d} -> (v, d)
335 Nothing -> (Strict.Nothing, empty) in
336 let fx = f cv in
337 let gm = go f p cm in
338 case (fx, size gm) of
339 (Strict.Nothing, 0) -> Nothing
340 (_, s) -> Just
341 Node
342 { node_value = fx
343 , node_descendants = gm
344 , node_size = s + 1
345 }
346 ) k m
347
348 -- * Fold
349
350 -- | Return the given accumulator folded by the given function
351 -- applied on non-'Strict.Nothing' 'node_value's
352 -- from left to right through the given 'TreeMap'.
353 foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
354 foldl_with_Path =
355 foldp []
356 where
357 foldp :: Ord k
358 => [k] -> (a -> Path k -> x -> a)
359 -> a -> TreeMap k x -> a
360 foldp p fct a (TreeMap m) =
361 Map.foldlWithKey
362 (\acc k Node{node_value, node_descendants} ->
363 let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
364 foldp (k:p) fct acc' node_descendants) a m
365
366 -- | Return the given accumulator folded by the given function
367 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
368 -- from left to right through the given 'TreeMap'.
369 foldl_with_Path_and_Node :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
370 foldl_with_Path_and_Node =
371 foldp []
372 where
373 foldp :: Ord k
374 => [k] -> (a -> Node k x -> Path k -> x -> a)
375 -> a -> TreeMap k x -> a
376 foldp p fct a (TreeMap m) =
377 Map.foldlWithKey
378 (\acc k n@Node{node_value, node_descendants} ->
379 let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
380 foldp (k:p) fct acc' node_descendants) a m
381
382 -- | Return the given accumulator folded by the given function
383 -- applied on non-'Strict.Nothing' 'node_value's
384 -- from right to left through the given 'TreeMap'.
385 foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
386 foldr_with_Path =
387 foldp []
388 where
389 foldp :: Ord k
390 => [k] -> (Path k -> x -> a -> a)
391 -> a -> TreeMap k x -> a
392 foldp p fct a (TreeMap m) =
393 Map.foldrWithKey
394 (\k Node{node_value, node_descendants} acc ->
395 let acc' = foldp (k:p) fct acc node_descendants in
396 Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
397
398 -- | Return the given accumulator folded by the given function
399 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
400 -- from right to left through the given 'TreeMap'.
401 foldr_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
402 foldr_with_Path_and_Node =
403 foldp []
404 where
405 foldp :: Ord k
406 => [k] -> (Node k x -> Path k -> x -> a -> a)
407 -> a -> TreeMap k x -> a
408 foldp p fct a (TreeMap m) =
409 Map.foldrWithKey
410 (\k n@Node{node_value, node_descendants} acc ->
411 let acc' = foldp (k:p) fct acc node_descendants in
412 Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
413
414 -- | Return the given accumulator folded by the given function
415 -- applied on non-'Strict.Nothing' 'node_value's
416 -- from left to right along the given 'Path'.
417 foldl_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
418 foldl_path fct =
419 go fct [] . list
420 where
421 go :: Ord k
422 => (Path k -> x -> a -> a) -> [k] -> [k]
423 -> TreeMap k x -> a -> a
424 go _f _ [] _t a = a
425 go f p (k:n) (TreeMap t) a =
426 case Map.lookup k t of
427 Nothing -> a
428 Just Node{node_value=v, node_descendants=d} ->
429 case v of
430 Strict.Nothing -> go f (k:p) n d a
431 Strict.Just x -> go f (k:p) n d (f (reverse $ path k p) x a)
432
433 -- | Return the given accumulator folded by the given function
434 -- applied on non-'Strict.Nothing' 'node_value's
435 -- from right to left along the given 'Path'.
436 foldr_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
437 foldr_path fct =
438 go fct [] . list
439 where
440 go :: Ord k
441 => (Path k -> x -> a -> a) -> [k] -> [k]
442 -> TreeMap k x -> a -> a
443 go _f _ [] _t a = a
444 go f p (k:n) (TreeMap t) a =
445 case Map.lookup k t of
446 Nothing -> a
447 Just Node{node_value=v, node_descendants=d} ->
448 case v of
449 Strict.Nothing -> go f (k:p) n d a
450 Strict.Just x -> f (reverse $ path k p) x $ go f (k:p) n d a
451
452 -- * Flatten
453
454 -- | Return a 'Map' associating each 'Path'
455 -- leading to a non-'Strict.Nothing' 'node_value' in the given 'TreeMap',
456 -- with its value mapped by the given function.
457 flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
458 flatten = flatten_with_Path . const
459
460 -- | Like 'flatten' but with also the current 'Path' given to the mapping function.
461 flatten_with_Path :: Ord k => (Path k -> x -> y) -> TreeMap k x -> Map (Path k) y
462 flatten_with_Path =
463 flat_map []
464 where
465 flat_map :: Ord k
466 => [k] -> (Path k -> x -> y)
467 -> TreeMap k x
468 -> Map (Path k) y
469 flat_map p f (TreeMap m) =
470 Map.unions $
471 Map.mapKeysMonotonic (reverse . flip path p) (
472 Map.mapMaybeWithKey (\k Node{node_value} ->
473 case node_value of
474 Strict.Nothing -> Nothing
475 Strict.Just x -> Just $ f (reverse $ path k p) x) m
476 ) :
477 Map.foldrWithKey
478 (\k -> (:) . flat_map (k:p) f . node_descendants)
479 [] m
480
481 -- * Filter
482
483 -- | Return the given 'TreeMap'
484 -- keeping only its non-'Strict.Nothing' 'node_value's
485 -- passing the given predicate.
486 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
487 filter f =
488 map_Maybe_with_Path
489 (\_p x -> if f x then Strict.Just x else Strict.Nothing)
490
491 -- | Like 'filter' but with also the current 'Path' given to the predicate.
492 filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
493 filter_with_Path f =
494 map_Maybe_with_Path
495 (\p x -> if f p x then Strict.Just x else Strict.Nothing)
496
497 -- | Like 'filter_with_Path' but with also the current 'Node' given to the predicate.
498 filter_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
499 filter_with_Path_and_Node f =
500 map_Maybe_with_Path_and_Node
501 (\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
502
503 -- | Return the given 'TreeMap'
504 -- mapping its non-'Strict.Nothing' 'node_value's
505 -- and keeping only the non-'Strict.Nothing' results.
506 map_Maybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
507 map_Maybe = map_Maybe_with_Path . const
508
509 -- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
510 map_Maybe_with_Path :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
511 map_Maybe_with_Path = map_Maybe_with_Path_and_Node . const
512
513 -- | Like 'map_Maybe_with_Path' but with also the current 'Node' given to the predicate.
514 map_Maybe_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
515 map_Maybe_with_Path_and_Node =
516 go []
517 where
518 go :: Ord k
519 => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
520 -> TreeMap k x
521 -> TreeMap k y
522 go p test (TreeMap m) =
523 TreeMap $
524 Map.mapMaybeWithKey
525 (\k node@Node{node_value=v, node_descendants=ns} ->
526 let node_descendants = go (k:p) test ns in
527 let node_size = size node_descendants in
528 case v of
529 Strict.Just x ->
530 let node_value = test node (reverse $ path k p) x in
531 case node_value of
532 Strict.Nothing | null node_descendants -> Nothing
533 Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
534 Strict.Just _ -> Just Node{node_value, node_descendants, node_size}
535 _ ->
536 if null node_descendants
537 then Nothing
538 else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
539 ) m