lib/Hcompta/Lib/TreeMap.hs

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