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, 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 -- * Union
 200
 201 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
 202 -- merging values (in respective order) when a 'Path' leads
 203 -- to a non-'Strict.Nothing' 'node_value' in both given 'TreeMap's.
 204 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
 205 union merge (TreeMap tm0) (TreeMap tm1) =
 206         TreeMap $
 207         Data.Map.unionWith
 208          (\Node{node_value=x0, node_descendants=m0}
 209            Node{node_value=x1, node_descendants=m1} ->
 210                 let m = union merge m0 m1 in
 211                 let x = Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0 in
 212                 Node
 213                  { node_value = x
 214                  , node_descendants = m
 215                  , node_size = size m + Strict.maybe 0 (const 1) x
 216                  })
 217          tm0 tm1
 218
 219 -- | Return the 'union' of the given 'TreeMap's.
 220 --
 221 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
 222 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 223 unions merge = Data.List.foldl' (union merge) empty
 224
 225 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 226 -- foldl' f = go
 227 --      where
 228 --              go z []     = z
 229 --              go z (x:xs) = z `seq` go (f z x) xs
 230
 231 -- * Map
 232
 233 -- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
 234 -- mapped by the given function.
 235 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 236 map f =
 237         TreeMap .
 238         Data.Map.map
 239          (\n@Node{node_value=x, node_descendants=m} ->
 240                 n{ node_value=Strict.maybe Strict.Nothing (Strict.Just . f) x
 241                  , node_descendants=Hcompta.Lib.TreeMap.map f m
 242                  }) .
 243         nodes
 244
 245 -- | Return the given 'TreeMap' with each 'node_value'
 246 -- mapped by the given function supplied with
 247 -- the already mapped 'node_descendants' of the current 'Node'.
 248 map_by_depth_first :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
 249 map_by_depth_first f =
 250         TreeMap .
 251         Data.Map.map
 252          (\Node{node_value, node_descendants} ->
 253                 let m = map_by_depth_first f node_descendants in
 254                 Node
 255                  { node_value = Strict.Just $ f m node_value
 256                  , node_descendants = m
 257                  , node_size = size m + 1
 258                  }) .
 259         nodes
 260
 261 -- * Alter
 262
 263 alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
 264 alterl_path fct =
 265         go fct . list
 266         where
 267                 go :: Ord k
 268                  => (Strict.Maybe x -> Strict.Maybe x) -> [k]
 269                  -> TreeMap k x -> TreeMap k x
 270                 go _f [] m = m
 271                 go f (k:p) (TreeMap m) =
 272                         TreeMap $
 273                         Data.Map.alter
 274                          (\c ->
 275                                 let (cv, cm) =
 276                                         case c of
 277                                          Just Node{node_value=v, node_descendants=d} -> (v, d)
 278                                          Nothing -> (Strict.Nothing, empty) in
 279                                 let fx = f cv in
 280                                 let gm = go f p cm in
 281                                 case (fx, size gm) of
 282                                  (Strict.Nothing, 0) -> Nothing
 283                                  (_, s) -> Just
 284                                         Node
 285                                          { node_value = fx
 286                                          , node_descendants = gm
 287                                          , node_size = s + 1
 288                                          }
 289                          ) k m
 290
 291 -- * Fold
 292
 293 -- | Return the given accumulator folded by the given function
 294 -- applied on non-'Strict.Nothing' 'node_value's
 295 -- from left to right through the given 'TreeMap'.
 296 foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 297 foldl_with_Path =
 298         foldp []
 299         where
 300                 foldp :: Ord k
 301                  => [k] -> (a -> Path k -> x -> a)
 302                  -> a -> TreeMap k x -> a
 303                 foldp p fct a (TreeMap m) =
 304                         Data.Map.foldlWithKey
 305                          (\acc k Node{node_value, node_descendants} ->
 306                                 let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
 307                                 foldp (k:p) fct acc' node_descendants) a m
 308
 309 -- | Return the given accumulator folded by the given function
 310 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
 311 -- from left to right through the given 'TreeMap'.
 312 foldl_with_Path_and_Node :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
 313 foldl_with_Path_and_Node =
 314         foldp []
 315         where
 316                 foldp :: Ord k
 317                  => [k] -> (a -> Node k x -> Path k -> x -> a)
 318                  -> a -> TreeMap k x -> a
 319                 foldp p fct a (TreeMap m) =
 320                         Data.Map.foldlWithKey
 321                          (\acc k n@Node{node_value, node_descendants} ->
 322                                 let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
 323                                 foldp (k:p) fct acc' node_descendants) a m
 324
 325 -- | Return the given accumulator folded by the given function
 326 -- applied on non-'Strict.Nothing' 'node_value's
 327 -- from right to left through the given 'TreeMap'.
 328 foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 329 foldr_with_Path =
 330         foldp []
 331         where
 332                 foldp :: Ord k
 333                  => [k] -> (Path k -> x -> a -> a)
 334                  -> a -> TreeMap k x -> a
 335                 foldp p fct a (TreeMap m) =
 336                         Data.Map.foldrWithKey
 337                          (\k Node{node_value, node_descendants} acc ->
 338                                 let acc' = foldp (k:p) fct acc node_descendants in
 339                                 Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) 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 right to left through the given 'TreeMap'.
 344 foldr_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 345 foldr_with_Path_and_Node =
 346         foldp []
 347         where
 348                 foldp :: Ord k
 349                  => [k] -> (Node k x -> Path k -> x -> a -> a)
 350                  -> a -> TreeMap k x -> a
 351                 foldp p fct a (TreeMap m) =
 352                         Data.Map.foldrWithKey
 353                          (\k n@Node{node_value, node_descendants} acc ->
 354                                 let acc' = foldp (k:p) fct acc node_descendants in
 355                                 Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
 356
 357 -- | Return the given accumulator folded by the given function
 358 -- applied on non-'Strict.Nothing' 'node_value's
 359 -- from left to right along the given 'Path'.
 360 foldl_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 361 foldl_path fct =
 362         go fct [] . list
 363         where
 364                 go :: Ord k
 365                  => (Path k -> x -> a -> a) -> [k] -> [k]
 366                  -> TreeMap k x -> a -> a
 367                 go _f _ [] _t a = a
 368                 go f p (k:n) (TreeMap t) a =
 369                         case Data.Map.lookup k t of
 370                          Nothing -> a
 371                          Just Node{node_value=v, node_descendants=d} ->
 372                                 case v of
 373                                  Strict.Nothing -> go f (k:p) n d a
 374                                  Strict.Just x  -> go f (k:p) n d (f (reverse $ path k p) x a)
 375
 376 -- | Return the given accumulator folded by the given function
 377 -- applied on non-'Strict.Nothing' 'node_value's
 378 -- from right to left along the given 'Path'.
 379 foldr_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 380 foldr_path fct =
 381         go fct [] . list
 382         where
 383                 go :: Ord k
 384                  => (Path k -> x -> a -> a) -> [k] -> [k]
 385                  -> TreeMap k x -> a -> a
 386                 go _f _ [] _t a = a
 387                 go f p (k:n) (TreeMap t) a =
 388                         case Data.Map.lookup k t of
 389                          Nothing -> a
 390                          Just Node{node_value=v, node_descendants=d} ->
 391                                 case v of
 392                                  Strict.Nothing -> go f (k:p) n d a
 393                                  Strict.Just x  -> f (reverse $ path k p) x $ go f (k:p) n d a
 394
 395 -- * Flatten
 396
 397 -- | Return a 'Map' associating each 'Path'
 398 -- leading to a non-'Strict.Nothing' 'node_value' in the given 'TreeMap',
 399 -- with its value mapped by the given function.
 400 flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
 401 flatten = flatten_with_Path . const
 402
 403 -- | Like 'flatten' but with also the current 'Path' given to the mapping function.
 404 flatten_with_Path :: Ord k => (Path k -> x -> y) -> TreeMap k x -> Map (Path k) y
 405 flatten_with_Path =
 406         flat_map []
 407         where
 408                 flat_map :: Ord k
 409                  => [k] -> (Path k -> x -> y)
 410                  -> TreeMap k x
 411                  -> Map (Path k) y
 412                 flat_map p f (TreeMap m) =
 413                         Data.Map.unions $
 414                         (
 415                         Data.Map.mapKeysMonotonic (reverse . flip path p) $
 416                         Data.Map.mapMaybeWithKey (\k Node{node_value} ->
 417                                 case node_value of
 418                                  Strict.Nothing -> Nothing
 419                                  Strict.Just x  -> Just $ f (reverse $ path k p) x) m
 420                         ) :
 421                         Data.Map.foldrWithKey
 422                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 423                          [] m
 424
 425 -- * Filter
 426
 427 -- | Return the given 'TreeMap'
 428 --   keeping only its non-'Strict.Nothing' 'node_value's
 429 --   passing the given predicate.
 430 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
 431 filter f =
 432         map_Maybe_with_Path
 433          (\_p x -> if f x then Strict.Just x else Strict.Nothing)
 434
 435 -- | Like 'filter' but with also the current 'Path' given to the predicate.
 436 filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 437 filter_with_Path f =
 438         map_Maybe_with_Path
 439          (\p x -> if f p x then Strict.Just x else Strict.Nothing)
 440
 441 -- | Like 'filter_with_Path' but with also the current 'Node' given to the predicate.
 442 filter_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 443 filter_with_Path_and_Node f =
 444         map_Maybe_with_Path_and_Node
 445          (\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
 446
 447 -- | Return the given 'TreeMap'
 448 --   mapping its non-'Strict.Nothing' 'node_value's
 449 --   and keeping only the non-'Strict.Nothing' results.
 450 map_Maybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 451 map_Maybe = map_Maybe_with_Path . const
 452
 453 -- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
 454 map_Maybe_with_Path :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 455 map_Maybe_with_Path = map_Maybe_with_Path_and_Node . const
 456
 457 -- | Like 'map_Maybe_with_Path' but with also the current 'Node' given to the predicate.
 458 map_Maybe_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 459 map_Maybe_with_Path_and_Node =
 460         go []
 461         where
 462                 go :: Ord k
 463                  => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
 464                  -> TreeMap k x
 465                  -> TreeMap k y
 466                 go p test (TreeMap m) =
 467                         TreeMap $
 468                         Data.Map.mapMaybeWithKey
 469                          (\k node@Node{node_value=v, node_descendants=ns} ->
 470                                 let node_descendants = go (k:p) test ns in
 471                                 let node_size = size node_descendants in
 472                                 case v of
 473                                  Strict.Just x ->
 474                                         let node_value = test node (reverse $ path k p) x in
 475                                         case node_value of
 476                                          Strict.Nothing | null node_descendants -> Nothing
 477                                          Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
 478                                          Strict.Just _  -> Just Node{node_value, node_descendants, node_size}
 479                                  _ ->
 480                                         if null node_descendants
 481                                         then Nothing
 482                                         else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
 483                          ) m