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