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