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 -> Path k -> Node k x -> x -> a) -> a -> TreeMap k x -> a
 301 foldl_with_Path_and_Node =
 302         foldp []
 303         where
 304                 foldp :: Ord k
 305                  => [k] -> (a -> Path k -> Node k x -> 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 (reverse $ path k p) n) 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 => (Path k -> Node k x -> x -> a -> a) -> a -> TreeMap k x -> a
 333 foldr_with_Path_and_Node =
 334         foldp []
 335         where
 336                 foldp :: Ord k
 337                  => [k] -> (Path k -> Node k x -> 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 (reverse $ path k p) n 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 -- | Return the given 'TreeMap'
 423 --   mapping its non-'Nothing' 'node_value's
 424 --   and keeping only the non-'Nothing' results.
 425 map_Maybe :: Ord k => (x -> Maybe y) -> TreeMap k x -> TreeMap k y
 426 map_Maybe f = map_Maybe_with_Path (const f)
 427
 428 -- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
 429 map_Maybe_with_Path :: Ord k => (Path k -> x -> Maybe y) -> TreeMap k x -> TreeMap k y
 430 map_Maybe_with_Path =
 431         go []
 432         where
 433                 go :: Ord k
 434                  => [k] -> (Path k -> x -> Maybe y)
 435                  -> TreeMap k x
 436                  -> TreeMap k y
 437                 go p test (TreeMap m) =
 438                         TreeMap $
 439                         Data.Map.mapMaybeWithKey
 440                          (\k Node{node_value=v, node_descendants=ns} ->
 441                                 let node_descendants = go (k:p) test ns in
 442                                 let node_size = size node_descendants in
 443                                 case v of
 444                                  Just x ->
 445                                         let node_value = test (reverse $ path k p) x in
 446                                         case node_value of
 447                                          Nothing | null node_descendants -> Nothing
 448                                          Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
 449                                          Just _  -> Just Node{node_value, node_descendants, node_size}
 450                                  _ ->
 451                                         if null node_descendants
 452                                         then Nothing
 453                                         else Just Node{node_value=Nothing, node_descendants, node_size}
 454                          ) m