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