lib/Hcompta/Lib/TreeMap.hs

   1 {-# LANGUAGE DeriveDataTypeable #-}
   2 {-# LANGUAGE NamedFieldPuns #-}
   3
   4 -- | This module implements a 'TreeMap',
   5 -- which is like a 'Data.Map.Map'
   6 -- but whose key is now a 'NonEmpty' list of 'Data.Map.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
  18 import           Data.Monoid (Monoid(..))
  19 import           Data.Traversable (Traversable(..))
  20 import           Data.Typeable (Typeable)
  21 import           Prelude hiding (reverse)
  22
  23 -- * The 'TreeMap' type
  24
  25 newtype TreeMap k x
  26  =      TreeMap (Data.Map.Map k (Node k x))
  27  deriving (Data, Eq, Read, Show, Typeable)
  28
  29 instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
  30         mempty = empty
  31         mappend = union const
  32         -- mconcat = Data.List.foldr mappend mempty
  33 instance Ord k => Functor (TreeMap k) where
  34         fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
  35 instance Ord k => Foldable (TreeMap k) where
  36         foldMap f (TreeMap m) = foldMap (foldMap f) m
  37 instance Ord k => Traversable (TreeMap k) where
  38         traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m
  39
  40 -- * The 'Path' type
  41
  42 -- | A 'Path' is a non-empty list of 'Data.Map.Map' keys.
  43 type Path k = NonEmpty k
  44
  45 path :: k -> [k] -> Path k
  46 path = (:|)
  47
  48 list :: Path k -> [k]
  49 list = Data.List.NonEmpty.toList
  50
  51 reverse :: Path k -> Path k
  52 reverse = Data.List.NonEmpty.reverse
  53
  54 -- * The 'Node' type
  55 data Ord k
  56  => Node k x
  57  =  Node
  58  { node_size        :: Int -- ^ The number of non-'Nothing' 'node_value's reachable from this 'Node'.
  59  , node_value       :: Maybe x -- ^ Some value, or 'Nothing' if this 'Node' is intermediary.
  60  , node_descendants :: TreeMap k x -- ^ Descendants 'Node's.
  61  } deriving (Data, Eq, Read, Show, Typeable)
  62
  63 instance (Ord k, Monoid v) => Monoid (Node k v) where
  64         mempty =
  65                 Node
  66                  { node_value = Nothing
  67                  , node_size = 0
  68                  , node_descendants = TreeMap mempty
  69                  }
  70         mappend
  71          Node{node_value=x0, node_descendants=m0}
  72          Node{node_value=x1, node_descendants=m1} =
  73                 let m = union const m0 m1 in
  74                 let x = x0 `mappend` x1 in
  75                 Node
  76                  { node_value = x
  77                  , node_size = size m + maybe 0 (const 1) x
  78                  , node_descendants = union const m0 m1
  79                  }
  80         -- mconcat = Data.List.foldr mappend mempty
  81
  82 instance Ord k => Functor (Node k) where
  83         fmap f Node{node_value=x, node_descendants=m, node_size} =
  84                 Node
  85                  { node_value = fmap f x
  86                  , node_descendants = Hcompta.Lib.TreeMap.map f m
  87                  , node_size
  88                  }
  89
  90 instance Ord k => Foldable (Node k) where
  91         foldMap f Node{node_value=Nothing, node_descendants=TreeMap m} =
  92                 foldMap (foldMap f) m
  93         foldMap f Node{node_value=Just x, node_descendants=TreeMap m} =
  94                 f x `mappend` foldMap (foldMap f) m
  95
  96 instance Ord k => Traversable (Node k) where
  97         traverse f Node{node_value=Nothing, node_descendants=TreeMap m, node_size} =
  98                 Node node_size <$> pure Nothing <*> (TreeMap <$> traverse (traverse f) m)
  99         traverse f Node{node_value=Just x, node_descendants=TreeMap m, node_size} =
 100                 Node node_size <$> (Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
 101
 102 -- * Contructors
 103
 104 -- | Return the empty 'TreeMap'.
 105 empty :: TreeMap k x
 106 empty = TreeMap Data.Map.empty
 107
 108 -- | Return a 'TreeMap' only mapping the given 'Path' to the given value.
 109 singleton :: Ord k => Path k -> x -> TreeMap k x
 110 singleton ks x = insert const ks x empty
 111
 112 -- | Return a 'Node' only containing the given value.
 113 leaf :: Ord k => x -> Node k x
 114 leaf x =
 115         Node
 116          { node_value = Just x
 117          , node_descendants = empty
 118          , node_size = 1
 119          }
 120
 121 -- | Return the given 'TreeMap' associating the given 'Path' with the given value,
 122 -- merging values if the given 'TreeMap' already associates the given 'Path'
 123 -- with a non-'Nothing' 'node_value'.
 124 insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
 125 insert merge (k:|[]) x (TreeMap m) =
 126         TreeMap $
 127         Data.Map.insertWith
 128          (\_ Node{node_value=x1, node_descendants=m1, node_size=s1} ->
 129                 Node
 130                  { node_value = maybe (Just x) (Just . merge x) x1
 131                  , node_descendants = m1
 132                  , node_size = maybe (s1 + 1) (const s1) x1
 133                  })
 134          k (leaf x) m
 135 insert merge (k:|k':ks) x (TreeMap m) =
 136         TreeMap $
 137         Data.Map.insertWith
 138          (\_ Node{node_value=x1, node_descendants=m1} ->
 139                 let m' = insert merge (path k' ks) x m1 in
 140                 Node{node_value=x1, node_descendants=m', node_size=size m' + maybe 0 (const 1) x1})
 141          k
 142          (Node
 143                  { node_value = Nothing
 144                  , node_descendants = insert merge (path k' ks) x empty
 145                  , node_size = 1
 146                  })
 147          m
 148
 149 -- | Return a 'TreeMap' associating for each tuple of the given list
 150 -- the 'Path' to the value,
 151 -- merging values of identical 'Path's (in respective order).
 152 from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
 153 from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
 154
 155 -- | Return a 'TreeMap' associating for each key and value of the given 'Data.Map.Map'
 156 -- the 'Path' to the value,
 157 -- merging values of identical 'Path's (in respective order).
 158 from_Map :: Ord k => (x -> x -> x) -> Data.Map.Map (Path k) x -> TreeMap k x
 159 from_Map merge = Data.Map.foldlWithKey (\acc p x -> insert merge p x acc) empty
 160
 161 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
 162 -- merging values (in respective order) when a 'Path' leads
 163 -- to a non-'Nothing' 'node_value' in both given 'TreeMap's.
 164 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
 165 union merge (TreeMap tm0) (TreeMap tm1) =
 166         TreeMap $
 167         Data.Map.unionWith
 168          (\Node{node_value=x0, node_descendants=m0}
 169            Node{node_value=x1, node_descendants=m1} ->
 170                 let m = union merge m0 m1 in
 171                 let x = maybe x1 (\x0' -> maybe (Just x0') (Just . merge x0') x1) x0 in
 172                 Node
 173                  { node_value = x
 174                  , node_descendants = m
 175                  , node_size = size m + maybe 0 (const 1) x
 176                  })
 177          tm0 tm1
 178
 179 -- | Return the 'union' of the given 'TreeMap's.
 180 --
 181 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
 182 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 183 unions merge ts = Data.List.foldl' (union merge) empty ts
 184
 185 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 186 -- foldl' f = go
 187 --      where
 188 --              go z []     = z
 189 --              go z (x:xs) = z `seq` go (f z x) xs
 190
 191 -- | Return the given 'TreeMap' with each non-'Nothing' 'node_value'
 192 -- mapped by the given function.
 193 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 194 map f =
 195         TreeMap .
 196         Data.Map.map
 197          (\n@Node{node_value=x, node_descendants=m} ->
 198                 n{ node_value=maybe Nothing (Just . f) x
 199                  , node_descendants=Hcompta.Lib.TreeMap.map f m
 200                  }) .
 201         nodes
 202
 203 -- | Return the given 'TreeMap' with each 'node_value'
 204 -- mapped by the given function supplied with
 205 -- the already mapped 'node_descendants' of the current 'Node'.
 206 map_by_depth_first :: Ord k => (TreeMap k y -> Maybe x -> y) -> TreeMap k x -> TreeMap k y
 207 map_by_depth_first f =
 208         TreeMap .
 209         Data.Map.map
 210          (\n@Node{node_value, node_descendants} ->
 211                 let m = map_by_depth_first f node_descendants in
 212                 let x = f m node_value in
 213                 n{ node_value = Just x
 214                  , node_descendants = m
 215                  , node_size = size m + 1
 216                  }) .
 217         nodes
 218
 219 -- * Extractors
 220
 221 -- | Return the 'Data.Map.Map' in the given 'TreeMap'.
 222 nodes :: TreeMap k x -> Data.Map.Map k (Node k x)
 223 nodes (TreeMap m) = m
 224
 225 -- | Return the number of non-'Nothing' 'node_value's in the given 'TreeMap'.
 226 --
 227 --   * Complexity: O(r) where r is the size of the root 'Data.Map.Map'.
 228 size :: Ord k => TreeMap k x -> Int
 229 size = Data.Map.foldr ((+) . node_size) 0 . nodes
 230
 231 -- | Return the value (if any) associated with the given 'Path'.
 232 find :: Ord k => Path k -> TreeMap k x -> Maybe x
 233 find (k:|[]) (TreeMap m) = maybe Nothing node_value $ Data.Map.lookup k m
 234 find (k:|k':ks) (TreeMap m) =
 235         maybe Nothing (find (path k' ks) . node_descendants) $
 236         Data.Map.lookup k m
 237
 238 -- | Return the given accumulator folded by the given function
 239 -- applied on non-'Nothing' 'node_value's
 240 -- from left to right through the given 'TreeMap'.
 241 foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 242 foldl_with_Path =
 243         foldp []
 244         where
 245                 foldp :: Ord k
 246                  => [k] -> (a -> Path k -> x -> a)
 247                  -> a -> TreeMap k x -> a
 248                 foldp p fct a (TreeMap m) =
 249                         Data.Map.foldlWithKey
 250                          (\acc k Node{node_value, node_descendants} ->
 251                                 let p' = path k p in
 252                                 let acc' = maybe acc (fct acc (reverse p')) node_value in
 253                                 foldp (k:p) fct acc' node_descendants) a m
 254
 255 -- | Return the given accumulator folded by the given function
 256 -- applied on non-'Nothing' 'Node's and 'node_value's
 257 -- from left to right through the given 'TreeMap'.
 258 foldl_with_Path_and_Node :: Ord k => (a -> Path k -> Node k x -> x -> a) -> a -> TreeMap k x -> a
 259 foldl_with_Path_and_Node =
 260         foldp []
 261         where
 262                 foldp :: Ord k
 263                  => [k] -> (a -> Path k -> Node k x -> x -> a)
 264                  -> a -> TreeMap k x -> a
 265                 foldp p fct a (TreeMap m) =
 266                         Data.Map.foldlWithKey
 267                          (\acc k n@Node{node_value, node_descendants} ->
 268                                 let p' = path k p in
 269                                 let acc' = maybe acc (fct acc (reverse p') n) node_value in
 270                                 foldp (k:p) fct acc' node_descendants) a m
 271
 272 -- | Return the given accumulator folded by the given function
 273 -- applied on non-'Nothing' 'node_value's
 274 -- from right to left through the given 'TreeMap'.
 275 foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 276 foldr_with_Path =
 277         foldp []
 278         where
 279                 foldp :: Ord k
 280                  => [k] -> (Path k -> x -> a -> a)
 281                  -> a -> TreeMap k x -> a
 282                 foldp p fct a (TreeMap m) =
 283                         Data.Map.foldrWithKey
 284                          (\k Node{node_value, node_descendants} acc ->
 285                                 let p' = path k p in
 286                                 let acc' = foldp (k:p) fct acc node_descendants in
 287                                 maybe acc' (\x -> fct (reverse p') x acc') node_value) a m
 288
 289 -- | Return the given accumulator folded by the given function
 290 -- applied on non-'Nothing' 'Node's and 'node_value's
 291 -- from right to left through the given 'TreeMap'.
 292 foldr_with_Path_and_Node :: Ord k => (Path k -> Node k x -> x -> a -> a) -> a -> TreeMap k x -> a
 293 foldr_with_Path_and_Node =
 294         foldp []
 295         where
 296                 foldp :: Ord k
 297                  => [k] -> (Path k -> Node k x -> x -> a -> a)
 298                  -> a -> TreeMap k x -> a
 299                 foldp p fct a (TreeMap m) =
 300                         Data.Map.foldrWithKey
 301                          (\k n@Node{node_value, node_descendants} acc ->
 302                                 let p' = path k p in
 303                                 let acc' = foldp (k:p) fct acc node_descendants in
 304                                 maybe acc' (\x -> fct (reverse p') n x acc') node_value) a m
 305
 306 -- | Return a 'Data.Map.Map' associating each 'Path'
 307 -- leading to a non-'Nothing' 'node_value' in the given 'TreeMap',
 308 -- with its value mapped by the given function.
 309 flatten :: Ord k => (x -> y) -> TreeMap k x -> Data.Map.Map (Path k) y
 310 flatten =
 311         flat_map []
 312         where
 313                 flat_map :: Ord k
 314                  => [k] -> (x -> y)
 315                  -> TreeMap k x
 316                  -> Data.Map.Map (Path k) y
 317                 flat_map p f (TreeMap m) =
 318                         Data.Map.unions $
 319                         (
 320                         Data.Map.mapKeysMonotonic (reverse . flip path p) $
 321                         Data.Map.mapMaybe (\Node{node_value=x} -> f <$> x) m
 322                         ) :
 323                         Data.Map.foldrWithKey
 324                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 325                          [] m