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