1 {-# LANGUAGE DeriveDataTypeable #-}
2 {-# LANGUAGE NamedFieldPuns #-}
4 -- | This module implements a strict 'TreeMap',
5 -- which is like a 'Map'
6 -- but whose key is now a 'NonEmpty' list of 'Map' keys (a 'Path')
7 -- enabling the possibility to gather mapped values
8 -- by 'Path' prefixes (inside a 'Node').
9 module Hcompta.Lib.TreeMap where
11 import Control.Applicative ((<$>), (<*>), pure)
12 import Control.DeepSeq (NFData(..))
15 import Data.Data (Data)
16 import Data.Foldable (Foldable, foldMap)
17 import Data.Functor (Functor(..))
18 import Data.Ord (Ord(..))
19 import qualified Data.List
20 import qualified Data.List.NonEmpty
21 import Data.List.NonEmpty (NonEmpty(..))
22 import Data.Map.Strict (Map)
23 import qualified Data.Map.Strict as Data.Map
24 import Data.Maybe (Maybe(..), maybe)
25 import Data.Monoid (Monoid(..))
26 import qualified Data.Strict.Maybe as Strict
27 import Data.Traversable (Traversable(..))
28 import Data.Typeable (Typeable)
29 import Prelude (($), (.), Int, Num(..), Show, const, flip, id, seq)
31 import qualified Hcompta.Lib.Strict as Strict ()
37 = TreeMap (Map k (Node k x))
38 deriving (Data, Eq, Show, Typeable)
40 instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
42 mappend = union mappend
43 -- mconcat = Data.List.foldr mappend mempty
44 instance Ord k => Functor (TreeMap k) where
45 fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
46 instance Ord k => Foldable (TreeMap k) where
47 foldMap f (TreeMap m) = foldMap (foldMap f) m
48 instance Ord k => Traversable (TreeMap k) where
49 traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m
50 instance (Ord k, NFData k, NFData x) => NFData (TreeMap k x) where
51 rnf (TreeMap m) = rnf m
55 -- | A 'Path' is a non-empty list of 'Map' keys.
56 type Path k = NonEmpty k
58 path :: k -> [k] -> Path k
62 list = Data.List.NonEmpty.toList
64 reverse :: Path k -> Path k
65 reverse = Data.List.NonEmpty.reverse
71 { node_size :: !Int -- ^ The number of non-'Strict.Nothing' 'node_value's reachable from this 'Node'.
72 , node_value :: !(Strict.Maybe x) -- ^ Some value, or 'Strict.Nothing' if this 'Node' is intermediary.
73 , node_descendants :: !(TreeMap k x) -- ^ Descendants 'Node's.
74 } deriving (Data, Eq, Show, Typeable)
77 instance (Ord k, Monoid v) => Monoid (Node k v) where
80 { node_value = Strict.Nothing
82 , node_descendants = TreeMap mempty
85 Node{node_value=x0, node_descendants=m0}
86 Node{node_value=x1, node_descendants=m1} =
87 let node_descendants = union const m0 m1 in
88 let node_value = x0 `mappend` x1 in
91 , node_size = size node_descendants
92 + Strict.maybe 0 (const 1) node_value
95 -- mconcat = Data.List.foldr mappend mempty
96 instance Ord k => Functor (Node k) where
97 fmap f Node{node_value=x, node_descendants=m, node_size} =
99 { node_value = fmap f x
100 , node_descendants = map f m
103 instance Ord k => Foldable (Node k) where
104 foldMap f Node{node_value=Strict.Nothing, node_descendants=TreeMap m} =
105 foldMap (foldMap f) m
106 foldMap f Node{node_value=Strict.Just x, node_descendants=TreeMap m} =
107 f x `mappend` foldMap (foldMap f) m
108 instance Ord k => Traversable (Node k) where
109 traverse f Node{node_value=Strict.Nothing, node_descendants=TreeMap m, node_size} =
110 Node node_size <$> pure Strict.Nothing <*> (TreeMap <$> traverse (traverse f) m)
111 traverse f Node{node_value=Strict.Just x, node_descendants=TreeMap m, node_size} =
112 Node node_size <$> (Strict.Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
113 instance (Ord k, NFData k, NFData x) => NFData (Node k x) where
114 rnf (Node s v d) = rnf s `seq` rnf v `seq` rnf d
116 node_find :: Ord k => [k] -> Node k x -> Strict.Maybe (Node k x)
117 node_find [] n = Strict.Just n
118 node_find (k:ks) (Node {node_descendants=TreeMap m}) =
119 maybe Strict.Nothing (node_find ks) $
124 -- | Return the empty 'TreeMap'.
125 empty :: Ord k => TreeMap k x
126 empty = TreeMap Data.Map.empty
128 -- | Return a 'TreeMap' only mapping the given 'Path' to the given value.
129 singleton :: Ord k => Path k -> x -> TreeMap k x
130 singleton ks x = insert const ks x empty
132 -- | Return a 'Node' only containing the given value.
133 leaf :: Ord k => x -> Node k x
136 { node_value = Strict.Just x
137 , node_descendants = empty
141 -- | Return the given 'TreeMap' associating the given 'Path' with the given value,
142 -- merging values if the given 'TreeMap' already associates the given 'Path'
143 -- with a non-'Strict.Nothing' 'node_value'.
144 insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
145 insert merge (k:|[]) x (TreeMap m) =
148 (\_ Node{node_value = x1, node_descendants = m1, node_size = s1} ->
150 { node_value = Strict.maybe (Strict.Just x) (Strict.Just . merge x) x1
151 , node_descendants = m1
152 , node_size = Strict.maybe (s1 + 1) (const s1) x1
155 insert merge (k:|k':ks) x (TreeMap m) =
158 (\_ Node{node_value = x1, node_descendants = m1} ->
159 let m' = insert merge (path k' ks) x $ m1 in
160 let s' = size m' + Strict.maybe 0 (const 1) x1 in
161 Node{node_value=x1, node_descendants=m', node_size=s'})
164 { node_value = Strict.Nothing
165 , node_descendants = insert merge (path k' ks) x empty
170 -- | Return a 'TreeMap' associating for each tuple of the given list
171 -- the 'Path' to the value,
172 -- merging values of identical 'Path's (in respective order).
173 from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
174 from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
176 -- | Return a 'TreeMap' associating for each key and value of the given 'Map'
177 -- the 'Path' to the value,
178 -- merging values of identical 'Path's (in respective order).
179 from_Map :: Ord k => (x -> x -> x) -> Map (Path k) x -> TreeMap k x
180 from_Map merge = Data.Map.foldlWithKey (\acc p x -> insert merge p x acc) empty
184 -- | Return the 'Map' in the given 'TreeMap'.
185 nodes :: Ord k => TreeMap k x -> Map k (Node k x)
186 nodes (TreeMap m) = m
188 -- | Return 'True' iif. the given 'TreeMap' is 'empty'.
189 null :: Ord k => TreeMap k x -> Bool
190 null (TreeMap m) = Data.Map.null m
192 -- | Return the number of non-'Strict.Nothing' 'node_value's in the given 'TreeMap'.
194 -- * Complexity: O(r) where r is the size of the root 'Map'.
195 size :: Ord k => TreeMap k x -> Int
196 size = Data.Map.foldr ((+) . node_size) 0 . nodes
200 -- | Return the value (if any) associated with the given 'Path'.
201 find :: Ord k => Path k -> TreeMap k x -> Strict.Maybe x
202 find (k:|[]) (TreeMap m) = maybe Strict.Nothing node_value $ Data.Map.lookup k m
203 find (k:|k':ks) (TreeMap m) =
204 maybe Strict.Nothing (find (path k' ks) . node_descendants) $
207 -- | Return the values (if any) associated with the prefixes of the given 'Path' (included).
208 find_along :: Ord k => Path k -> TreeMap k x -> [x]
209 find_along p (TreeMap tm) =
212 go :: Ord k => [k] -> Map k (Node k x) -> [x]
215 case Data.Map.lookup k m of
218 Strict.maybe id (:) (node_value node) $
219 go ks $ nodes (node_descendants node)
221 find_node :: Ord k => Path k -> TreeMap k x -> Strict.Maybe (Node k x)
222 find_node (k:|[]) (TreeMap m) = maybe Strict.Nothing Strict.Just $ Data.Map.lookup k m
223 find_node (k:|k':ks) (TreeMap m) =
224 maybe Strict.Nothing (find_node (path k' ks) . node_descendants) $
229 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
230 -- merging values (in respective order) when a 'Path' leads
231 -- to a non-'Strict.Nothing' 'node_value' in both given 'TreeMap's.
232 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
233 union merge (TreeMap tm0) (TreeMap tm1) =
236 (\Node{node_value=x0, node_descendants=m0}
237 Node{node_value=x1, node_descendants=m1} ->
238 let node_descendants = union merge m0 m1 in
239 let node_value = Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0 in
241 { node_size = size node_descendants + Strict.maybe 0 (const 1) node_value
247 -- | Return the 'union' of the given 'TreeMap's.
249 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
250 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
251 unions merge = Data.List.foldl' (union merge) empty
253 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
257 -- go z (x:xs) = z `seq` go (f z x) xs
261 -- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
262 -- mapped by the given function.
263 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
267 (\n@Node{node_value=x, node_descendants=m} ->
268 n{ node_value = fmap f x
269 , node_descendants = map f m
273 -- | Return the given 'TreeMap' with each 'Path' section
274 -- and each non-'Strict.Nothing' 'node_value'
275 -- mapped by the given functions.
277 -- WARNING: the function mapping 'Path' sections must be monotonic,
278 -- like in 'Data.Map.mapKeysMonotonic'.
279 map_monotonic :: (Ord k, Ord l) => (k -> l) -> (x -> y) -> TreeMap k x -> TreeMap l y
280 map_monotonic fk fx =
282 Data.Map.mapKeysMonotonic fk .
284 (\n@Node{node_value=x, node_descendants=m} ->
285 n{ node_value = fmap fx x
286 , node_descendants = map_monotonic fk fx m
290 -- | Return the given 'TreeMap' with each 'node_value'
291 -- mapped by the given function supplied with
292 -- the already mapped 'node_descendants' of the current 'Node'.
293 map_by_depth_first :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
294 map_by_depth_first f =
297 (\Node{node_value, node_descendants} ->
298 let m = map_by_depth_first f node_descendants in
300 { node_value = Strict.Just $ f m node_value
301 , node_descendants = m
302 , node_size = size m + 1
308 alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
313 => (Strict.Maybe x -> Strict.Maybe x) -> [k]
314 -> TreeMap k x -> TreeMap k x
316 go f (k:p) (TreeMap m) =
322 Just Node{node_value=v, node_descendants=d} -> (v, d)
323 Nothing -> (Strict.Nothing, empty) in
325 let gm = go f p cm in
326 case (fx, size gm) of
327 (Strict.Nothing, 0) -> Nothing
331 , node_descendants = gm
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
346 => [k] -> (a -> Path k -> x -> a)
347 -> a -> TreeMap k x -> a
348 foldp p fct a (TreeMap m) =
349 Data.Map.foldlWithKey
350 (\acc k Node{node_value, node_descendants} ->
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
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 =
362 => [k] -> (a -> Node k x -> Path k -> x -> a)
363 -> a -> TreeMap k x -> a
364 foldp p fct a (TreeMap m) =
365 Data.Map.foldlWithKey
366 (\acc k n@Node{node_value, node_descendants} ->
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
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
378 => [k] -> (Path k -> x -> a -> a)
379 -> a -> TreeMap k x -> a
380 foldp p fct a (TreeMap m) =
381 Data.Map.foldrWithKey
382 (\k Node{node_value, node_descendants} 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
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 =
394 => [k] -> (Node k x -> Path k -> x -> a -> a)
395 -> a -> TreeMap k x -> a
396 foldp p fct a (TreeMap m) =
397 Data.Map.foldrWithKey
398 (\k n@Node{node_value, node_descendants} 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
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
410 => (Path k -> x -> a -> a) -> [k] -> [k]
411 -> TreeMap k x -> a -> a
413 go f p (k:n) (TreeMap t) a =
414 case Data.Map.lookup k t of
416 Just Node{node_value=v, node_descendants=d} ->
418 Strict.Nothing -> go f (k:p) n d a
419 Strict.Just x -> go f (k:p) n d (f (reverse $ path k p) x a)
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
429 => (Path k -> x -> a -> a) -> [k] -> [k]
430 -> TreeMap k x -> a -> a
432 go f p (k:n) (TreeMap t) a =
433 case Data.Map.lookup k t of
435 Just Node{node_value=v, node_descendants=d} ->
437 Strict.Nothing -> go f (k:p) n d a
438 Strict.Just x -> f (reverse $ path k p) x $ go f (k:p) n d a
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
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
454 => [k] -> (Path k -> x -> y)
457 flat_map p f (TreeMap m) =
460 Data.Map.mapKeysMonotonic (reverse . flip path p) $
461 Data.Map.mapMaybeWithKey (\k Node{node_value} ->
463 Strict.Nothing -> Nothing
464 Strict.Just x -> Just $ f (reverse $ path k p) x) m
466 Data.Map.foldrWithKey
467 (\k -> (:) . flat_map (k:p) f . node_descendants)
472 -- | Return the given 'TreeMap'
473 -- keeping only its non-'Strict.Nothing' 'node_value's
474 -- passing the given predicate.
475 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
478 (\_p x -> if f x then Strict.Just x else Strict.Nothing)
480 -- | Like 'filter' but with also the current 'Path' given to the predicate.
481 filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
484 (\p x -> if f p x then Strict.Just x else Strict.Nothing)
486 -- | Like 'filter_with_Path' but with also the current 'Node' given to the predicate.
487 filter_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
488 filter_with_Path_and_Node f =
489 map_Maybe_with_Path_and_Node
490 (\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
492 -- | Return the given 'TreeMap'
493 -- mapping its non-'Strict.Nothing' 'node_value's
494 -- and keeping only the non-'Strict.Nothing' results.
495 map_Maybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
496 map_Maybe = map_Maybe_with_Path . const
498 -- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
499 map_Maybe_with_Path :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
500 map_Maybe_with_Path = map_Maybe_with_Path_and_Node . const
502 -- | Like 'map_Maybe_with_Path' but with also the current 'Node' given to the predicate.
503 map_Maybe_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
504 map_Maybe_with_Path_and_Node =
508 => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
511 go p test (TreeMap m) =
513 Data.Map.mapMaybeWithKey
514 (\k node@Node{node_value=v, node_descendants=ns} ->
515 let node_descendants = go (k:p) test ns in
516 let node_size = size node_descendants in
519 let node_value = test node (reverse $ path k p) x in
521 Strict.Nothing | null node_descendants -> Nothing
522 Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
523 Strict.Just _ -> Just Node{node_value, node_descendants, node_size}
525 if null node_descendants
527 else Just Node{node_value=Strict.Nothing, node_descendants, node_size}