lib/Hcompta/Lib/TreeMap.hs

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