lib/Hcompta/Lib/TreeMap.hs

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