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_content's reachable from this 'Node'.
  59  , node_content     :: Maybe x -- ^ Some content, 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_content = Nothing
  67                  , node_size = 0
  68                  , node_descendants = TreeMap mempty
  69                  }
  70         mappend
  71          Node{node_content=x0, node_descendants=m0}
  72          Node{node_content=x1, node_descendants=m1} =
  73                 let m = union const m0 m1 in
  74                 let x = x0 `mappend` x1 in
  75                 Node
  76                  { node_content = 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_content=x, node_descendants=m, node_size} =
  84                 Node
  85                  { node_content = 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_content=Nothing, node_descendants=TreeMap m} =
  92                 foldMap (foldMap f) m
  93         foldMap f Node{node_content=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_content=Nothing, node_descendants=TreeMap m, node_size} =
  98                 Node node_size <$> pure Nothing <*> (TreeMap <$> traverse (traverse f) m)
  99         traverse f Node{node_content=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_content     = Just x
 117          , node_descendants = empty
 118          , node_size        = 1
 119          }
 120
 121 -- | Return the given 'TreeMap' associating the given 'Path' with the given content,
 122 -- merging contents if the given 'TreeMap' already associates the given 'Path'
 123 -- with a non-'Nothing' 'node_content'.
 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_content=x1, node_descendants=m1, node_size=s1} ->
 129                 Node
 130                  { node_content = 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_content=x1, node_descendants=m1} ->
 139                 let m' = insert merge (path k' ks) x m1 in
 140                 Node{node_content=x1, node_descendants=m', node_size=size m' + maybe 0 (const 1) x1})
 141          k
 142          (Node
 143                  { node_content = Nothing
 144                  , node_descendants = insert merge (path k' ks) x empty
 145                  , node_size = 1
 146                  })
 147          m
 148
 149 -- | Return a 'TreeMap' associating the given 'Path' to the given content,
 150 -- merging content of identical 'Path's (in respective order).
 151 from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
 152 from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
 153
 154 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
 155 -- merging contents (in respective order) when a 'Path' leads
 156 -- to a non-'Nothing' 'node_content' in both given 'TreeMap's.
 157 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
 158 union merge (TreeMap tm0) (TreeMap tm1) =
 159         TreeMap $
 160         Data.Map.unionWith
 161          (\Node{node_content=x0, node_descendants=m0}
 162            Node{node_content=x1, node_descendants=m1} ->
 163                 let m = union merge m0 m1 in
 164                 let x = maybe x1 (\x0' -> maybe (Just x0') (Just . merge x0') x1) x0 in
 165                 Node
 166                  { node_content = x
 167                  , node_descendants = m
 168                  , node_size = size m + maybe 0 (const 1) x
 169                  })
 170          tm0 tm1
 171
 172 -- | Return the 'union' of the given 'TreeMap's.
 173 --
 174 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
 175 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 176 unions merge ts = Data.List.foldl' (union merge) empty ts
 177
 178 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 179 -- foldl' f = go
 180 --      where
 181 --              go z []     = z
 182 --              go z (x:xs) = z `seq` go (f z x) xs
 183
 184 -- | Return the given 'TreeMap' with each non-'Nothing' 'node_content'
 185 -- mapped by the given function.
 186 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 187 map f =
 188         TreeMap .
 189         Data.Map.map
 190          (\n@Node{node_content=x, node_descendants=m} ->
 191                 n{ node_content=maybe Nothing (Just . f) x
 192                  , node_descendants=Hcompta.Lib.TreeMap.map f m
 193                  }) .
 194         nodes
 195
 196 -- | Return the given 'TreeMap' with each 'node_content'
 197 -- mapped by the given function supplied with
 198 -- the already mapped 'node_descendants' of the current 'Node'.
 199 map_by_depth_first :: Ord k => (TreeMap k y -> Maybe x -> y) -> TreeMap k x -> TreeMap k y
 200 map_by_depth_first f =
 201         TreeMap .
 202         Data.Map.map
 203          (\n@Node{node_content, node_descendants} ->
 204                 let m = map_by_depth_first f node_descendants in
 205                 let x = f m node_content in
 206                 n{ node_content = Just x
 207                  , node_descendants = m
 208                  , node_size = size m + 1
 209                  }) .
 210         nodes
 211
 212 -- * Extractors
 213
 214 -- | Return the 'Data.Map.Map' in the given 'TreeMap'.
 215 nodes :: TreeMap k x -> Data.Map.Map k (Node k x)
 216 nodes (TreeMap m) = m
 217
 218 -- | Return the number of non-'Nothing' 'node_content's in the given 'TreeMap'.
 219 size :: Ord k => TreeMap k x -> Int
 220 size = Data.Map.foldr ((+) . node_size) 0 . nodes
 221
 222 -- | Return the content (if any) associated with the given 'Path'.
 223 find :: Ord k => Path k -> TreeMap k x -> Maybe x
 224 find (k:|[]) (TreeMap m) = maybe Nothing node_content $ Data.Map.lookup k m
 225 find (k:|k':ks) (TreeMap m) =
 226         maybe Nothing (find (path k' ks) . node_descendants) $
 227         Data.Map.lookup k m
 228
 229 -- | Return the given accumulator folded by the given function
 230 -- applied on non-'Nothing' 'node_content's
 231 -- from left to right through the given 'TreeMap'.
 232 foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 233 foldl_with_Path =
 234         foldp []
 235         where
 236                 foldp :: Ord k
 237                  => [k] -> (a -> Path k -> x -> a)
 238                  -> a -> TreeMap k x -> a
 239                 foldp p fct a (TreeMap m) =
 240                         Data.Map.foldlWithKey
 241                          (\acc k Node{node_content, node_descendants} ->
 242                                 let p' = path k p in
 243                                 let acc' = maybe acc (fct acc (reverse p')) node_content in
 244                                 foldp (k:p) fct acc' node_descendants) a m
 245
 246 -- | Return the given accumulator folded by the given function
 247 -- applied on non-'Nothing' 'Node's and 'node_content's
 248 -- from left to right through the given 'TreeMap'.
 249 foldl_with_Path_and_Node :: Ord k => (a -> Path k -> Node k x -> x -> a) -> a -> TreeMap k x -> a
 250 foldl_with_Path_and_Node =
 251         foldp []
 252         where
 253                 foldp :: Ord k
 254                  => [k] -> (a -> Path k -> Node k x -> x -> a)
 255                  -> a -> TreeMap k x -> a
 256                 foldp p fct a (TreeMap m) =
 257                         Data.Map.foldlWithKey
 258                          (\acc k n@Node{node_content, node_descendants} ->
 259                                 let p' = path k p in
 260                                 let acc' = maybe acc (fct acc (reverse p') n) node_content in
 261                                 foldp (k:p) fct acc' node_descendants) a m
 262
 263 -- | Return the given accumulator folded by the given function
 264 -- applied on non-'Nothing' 'node_content's
 265 -- from right to left through the given 'TreeMap'.
 266 foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 267 foldr_with_Path =
 268         foldp []
 269         where
 270                 foldp :: Ord k
 271                  => [k] -> (Path k -> x -> a -> a)
 272                  -> a -> TreeMap k x -> a
 273                 foldp p fct a (TreeMap m) =
 274                         Data.Map.foldrWithKey
 275                          (\k Node{node_content, node_descendants} acc ->
 276                                 let p' = path k p in
 277                                 let acc' = foldp (k:p) fct acc node_descendants in
 278                                 maybe acc' (\x -> fct (reverse p') x acc') node_content) a m
 279
 280 -- | Return the given accumulator folded by the given function
 281 -- applied on non-'Nothing' 'Node's and 'node_content's
 282 -- from right to left through the given 'TreeMap'.
 283 foldr_with_Path_and_Node :: Ord k => (Path k -> Node k x -> x -> a -> a) -> a -> TreeMap k x -> a
 284 foldr_with_Path_and_Node =
 285         foldp []
 286         where
 287                 foldp :: Ord k
 288                  => [k] -> (Path k -> Node k x -> x -> a -> a)
 289                  -> a -> TreeMap k x -> a
 290                 foldp p fct a (TreeMap m) =
 291                         Data.Map.foldrWithKey
 292                          (\k n@Node{node_content, node_descendants} acc ->
 293                                 let p' = path k p in
 294                                 let acc' = foldp (k:p) fct acc node_descendants in
 295                                 maybe acc' (\x -> fct (reverse p') n x acc') node_content) a m
 296
 297 -- | Return a 'Data.Map.Map' associating each 'Path'
 298 -- leading to a non-'Nothing' 'node_content' in the given 'TreeMap',
 299 -- with its content mapped by the given function.
 300 flatten :: Ord k => (x -> y) -> TreeMap k x -> Data.Map.Map (Path k) y
 301 flatten =
 302         flat_map []
 303         where
 304                 flat_map :: Ord k
 305                  => [k] -> (x -> y)
 306                  -> TreeMap k x
 307                  -> Data.Map.Map (Path k) y
 308                 flat_map p f (TreeMap m) =
 309                         Data.Map.unions $
 310                         (
 311                         Data.Map.mapKeysMonotonic (reverse . flip path p) $
 312                         Data.Map.mapMaybe (\Node{node_content=x} -> f <$> x) m
 313                         ) :
 314                         Data.Map.foldrWithKey
 315                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 316                          [] m