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 -- * Fold
 250
 251 -- | Return the given accumulator folded by the given function
 252 -- applied on non-'Nothing' 'node_value's
 253 -- from left to right through the given 'TreeMap'.
 254 foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 255 foldl_with_Path =
 256         foldp []
 257         where
 258                 foldp :: Ord k
 259                  => [k] -> (a -> Path k -> x -> a)
 260                  -> a -> TreeMap k x -> a
 261                 foldp p fct a (TreeMap m) =
 262                         Data.Map.foldlWithKey
 263                          (\acc k Node{node_value, node_descendants} ->
 264                                 let acc' = maybe acc (fct acc (reverse $ path k p)) node_value in
 265                                 foldp (k:p) fct acc' node_descendants) a m
 266
 267 -- | Return the given accumulator folded by the given function
 268 -- applied on non-'Nothing' 'Node's and 'node_value's
 269 -- from left to right through the given 'TreeMap'.
 270 foldl_with_Path_and_Node :: Ord k => (a -> Path k -> Node k x -> x -> a) -> a -> TreeMap k x -> a
 271 foldl_with_Path_and_Node =
 272         foldp []
 273         where
 274                 foldp :: Ord k
 275                  => [k] -> (a -> Path k -> Node k x -> x -> a)
 276                  -> a -> TreeMap k x -> a
 277                 foldp p fct a (TreeMap m) =
 278                         Data.Map.foldlWithKey
 279                          (\acc k n@Node{node_value, node_descendants} ->
 280                                 let acc' = maybe acc (fct acc (reverse $ path k p) n) node_value in
 281                                 foldp (k:p) fct acc' node_descendants) a m
 282
 283 -- | Return the given accumulator folded by the given function
 284 -- applied on non-'Nothing' 'node_value's
 285 -- from right to left through the given 'TreeMap'.
 286 foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 287 foldr_with_Path =
 288         foldp []
 289         where
 290                 foldp :: Ord k
 291                  => [k] -> (Path k -> x -> a -> a)
 292                  -> a -> TreeMap k x -> a
 293                 foldp p fct a (TreeMap m) =
 294                         Data.Map.foldrWithKey
 295                          (\k Node{node_value, node_descendants} acc ->
 296                                 let acc' = foldp (k:p) fct acc node_descendants in
 297                                 maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
 298
 299 -- | Return the given accumulator folded by the given function
 300 -- applied on non-'Nothing' 'Node's and 'node_value's
 301 -- from right to left through the given 'TreeMap'.
 302 foldr_with_Path_and_Node :: Ord k => (Path k -> Node k x -> x -> a -> a) -> a -> TreeMap k x -> a
 303 foldr_with_Path_and_Node =
 304         foldp []
 305         where
 306                 foldp :: Ord k
 307                  => [k] -> (Path k -> Node k x -> x -> a -> a)
 308                  -> a -> TreeMap k x -> a
 309                 foldp p fct a (TreeMap m) =
 310                         Data.Map.foldrWithKey
 311                          (\k n@Node{node_value, node_descendants} acc ->
 312                                 let acc' = foldp (k:p) fct acc node_descendants in
 313                                 maybe acc' (\x -> fct (reverse $ path k p) n x acc') node_value) a m
 314
 315 -- * Flatten
 316
 317 -- | Return a 'Map' associating each 'Path'
 318 -- leading to a non-'Nothing' 'node_value' in the given 'TreeMap',
 319 -- with its value mapped by the given function.
 320 flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
 321 flatten =
 322         flat_map []
 323         where
 324                 flat_map :: Ord k
 325                  => [k] -> (x -> y)
 326                  -> TreeMap k x
 327                  -> Map (Path k) y
 328                 flat_map p f (TreeMap m) =
 329                         Data.Map.unions $
 330                         (
 331                         Data.Map.mapKeysMonotonic (reverse . flip path p) $
 332                         Data.Map.mapMaybe (\Node{node_value=x} -> f <$> x) m
 333                         ) :
 334                         Data.Map.foldrWithKey
 335                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 336                          [] m
 337
 338 -- * Filter
 339
 340 -- | Return the given 'TreeMap'
 341 --   keeping only its non-'Nothing' 'node_value's
 342 --   passing the given predicate.
 343 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
 344 filter f =
 345         map_Maybe_with_Path
 346          (\_p x -> if f x then Just x else Nothing)
 347
 348 -- | Like 'filter' but with also the current 'Path' given to the predicate.
 349 filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 350 filter_with_Path f =
 351         map_Maybe_with_Path
 352          (\p x -> if f p x then Just x else Nothing)
 353
 354 -- | Return the given 'TreeMap'
 355 --   mapping its non-'Nothing' 'node_value's
 356 --   and keeping only the non-'Nothing' results.
 357 map_Maybe :: Ord k => (x -> Maybe y) -> TreeMap k x -> TreeMap k y
 358 map_Maybe f = map_Maybe_with_Path (const f)
 359
 360 -- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
 361 map_Maybe_with_Path :: Ord k => (Path k -> x -> Maybe y) -> TreeMap k x -> TreeMap k y
 362 map_Maybe_with_Path =
 363         go []
 364         where
 365                 go :: Ord k
 366                  => [k] -> (Path k -> x -> Maybe y)
 367                  -> TreeMap k x
 368                  -> TreeMap k y
 369                 go p test (TreeMap m) =
 370                         TreeMap $
 371                         Data.Map.mapMaybeWithKey
 372                          (\k Node{node_value=v, node_descendants=ns} ->
 373                                 let node_descendants = go (k:p) test ns in
 374                                 let node_size = size node_descendants in
 375                                 case v of
 376                                  Just x ->
 377                                         let node_value = test (reverse $ path k p) x in
 378                                         case node_value of
 379                                          Nothing | null node_descendants -> Nothing
 380                                          Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
 381                                          Just _  -> Just Node{node_value, node_descendants, node_size}
 382                                  _ ->
 383                                         if null node_descendants
 384                                         then Nothing
 385                                         else Just Node{node_value=Nothing, node_descendants, node_size}
 386                          ) m