lib/Hcompta/Lib/TreeMap.hs

   1 {-# LANGUAGE DeriveDataTypeable #-}
   2 {-# LANGUAGE NamedFieldPuns #-}
   3
   4 -- | This module implements a tree of 'Data.Map.Map'.
   5 module Hcompta.Lib.TreeMap where
   6
   7 import           Control.Applicative ((<$>), (<*>), pure)
   8 import           Data.Data (Data)
   9 import           Data.Foldable (Foldable(..))
  10 import qualified Data.List
  11 import qualified Data.List.NonEmpty
  12 import           Data.List.NonEmpty (NonEmpty(..))
  13 import qualified Data.Map
  14 import           Data.Monoid (Monoid(..))
  15 import           Data.Traversable (Traversable(..))
  16 import           Data.Typeable (Typeable)
  17
  18 -- * The 'Path' type
  19
  20 -- | A 'Path' is a non-empty list.
  21 type Path k = NonEmpty k
  22
  23 path :: k -> [k] -> Path k
  24 path = (:|)
  25
  26 list :: Path k -> [k]
  27 list = Data.List.NonEmpty.toList
  28
  29 rev :: Path k -> Path k
  30 rev = Data.List.NonEmpty.reverse
  31
  32 -- * The 'TreeMap' type
  33
  34 type TreeMap k x = Data.Map.Map k (Node k x)
  35 data Ord k => Node k x
  36  =   Node
  37  { node_size        :: Int -- ^ The number of non-'Nothing' 'node_content's reachable from this 'Node'.
  38  , node_content     :: Maybe x -- ^ Some content, or 'Nothing' if this 'Node' is intermediary.
  39  , node_descendants :: TreeMap k x -- ^ Descendants 'Node's.
  40  } deriving (Data, Eq, Read, Show, Typeable)
  41
  42 instance (Ord k, Monoid v) => Monoid (Node k v) where
  43         mempty =
  44                 Node
  45                  { node_content = Nothing
  46                  , node_size = 0
  47                  , node_descendants = mempty
  48                  }
  49         mappend
  50          Node{node_content=x0, node_descendants=m0}
  51          Node{node_content=x1, node_descendants=m1} =
  52                 let m = union const m0 m1 in
  53                 let x = x0 `mappend` x1 in
  54                 Node
  55                  { node_content = x
  56                  , node_size = size m + maybe 0 (const 1) x
  57                  , node_descendants = union const m0 m1
  58                  }
  59         -- mconcat = Data.List.foldr mappend mempty
  60
  61 instance Ord k => Functor (Node k) where
  62         fmap f Node{node_content=x, node_descendants=m, node_size} =
  63                 Node
  64                  { node_content = fmap f x
  65                  , node_descendants = Hcompta.Lib.TreeMap.map f m
  66                  , node_size
  67                  }
  68
  69 instance Ord k => Foldable (Node k) where
  70         foldMap f Node{node_content=Nothing, node_descendants=m} =
  71                 foldMap (foldMap f) m
  72         foldMap f Node{node_content=Just x, node_descendants=m} =
  73                 f x `mappend` foldMap (foldMap f) m
  74
  75 instance Ord k => Traversable (Node k) where
  76         traverse f Node{node_content=Nothing, node_descendants=m, node_size} =
  77                 Node node_size <$> pure Nothing <*> traverse (traverse f) m
  78         traverse f Node{node_content=Just x, node_descendants=m, node_size} =
  79                 Node node_size <$> (Just <$> f x) <*> traverse (traverse f) m
  80
  81 -- * Contructors
  82
  83 empty :: TreeMap k x
  84 empty = Data.Map.empty
  85
  86 singleton :: Ord k => Path k -> x -> TreeMap k x
  87 singleton ks x = insert const ks x Data.Map.empty
  88
  89 leaf :: Ord k => x -> Node k x
  90 leaf x =
  91         Node
  92          { node_content     = Just x
  93          , node_descendants = Data.Map.empty
  94          , node_size        = 1
  95          }
  96
  97 -- | Return the given 'TreeMap' associating the given 'Path' with the given content,
  98 -- merging contents if the given 'TreeMap' already associates the given 'Path'
  99 -- with a non-'Nothing' 'node_content'.
 100 insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
 101 insert merge (k:|[]) x m =
 102         Data.Map.insertWith
 103          (\_ Node{node_content=x1, node_descendants=m1, node_size=s1} ->
 104                 Node
 105                  { node_content = maybe (Just x) (Just . merge x) x1
 106                  , node_descendants = m1
 107                  , node_size = maybe (s1 + 1) (const s1) x1
 108                  })
 109          k (leaf x) m
 110 insert merge (k:|k':ks) x m =
 111         Data.Map.insertWith
 112          (\_ Node{node_content=x1, node_descendants=m1} ->
 113                 let m' = insert merge (path k' ks) x m1 in
 114                 Node{node_content=x1, node_descendants=m', node_size=size m' + maybe 0 (const 1) x1})
 115          k
 116          (Node
 117                  { node_content = Nothing
 118                  , node_descendants = insert merge (path k' ks) x Data.Map.empty
 119                  , node_size = 1
 120                  })
 121          m
 122
 123 -- | Return a 'TreeMap' associating the given 'Path' to the given content,
 124 -- merging content of identical 'Path's (in respective order).
 125 from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
 126 from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty
 127
 128 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
 129 -- merging contents (in respective order) when a 'Path' leads
 130 -- to a non-'Nothing' 'node_content' in both given 'TreeMap's.
 131 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
 132 union merge =
 133         Data.Map.unionWith
 134          (\Node{node_content=x0, node_descendants=m0}
 135            Node{node_content=x1, node_descendants=m1} ->
 136                 let m = union merge m0 m1 in
 137                 let x = maybe x1 (\x0' -> maybe (Just x0') (Just . merge x0') x1) x0 in
 138                 Node
 139                  { node_content = x
 140                  , node_descendants = m
 141                  , node_size = size m + maybe 0 (const 1) x
 142                  })
 143
 144 -- | Return the 'union' of the given 'TreeMap's.
 145 --
 146 -- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
 147 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 148 unions merge ts = Data.List.foldl' (union merge) empty ts
 149
 150 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 151 -- foldl' f = go
 152 --      where
 153 --              go z []     = z
 154 --              go z (x:xs) = z `seq` go (f z x) xs
 155
 156 -- | Return the given 'TreeMap' with each non-'Nothing' 'node_content'
 157 -- mapped by the given function.
 158 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 159 map f =
 160         Data.Map.map
 161          (\n@Node{node_content=x, node_descendants=m} ->
 162                 n{ node_content=maybe Nothing (Just . f) x
 163                  , node_descendants=Hcompta.Lib.TreeMap.map f m
 164                  })
 165
 166 -- | Return the given 'TreeMap' with each 'node_content'
 167 -- mapped by the given function supplied with
 168 -- the already mapped 'node_descendants' of the current 'Node'.
 169 depth_first_map :: Ord k => (TreeMap k y -> Maybe x -> y) -> TreeMap k x -> TreeMap k y
 170 depth_first_map f =
 171         Data.Map.map
 172          (\n@Node{node_content, node_descendants} ->
 173                 let m = depth_first_map f node_descendants in
 174                 let x = f m node_content in
 175                 n{ node_content = Just x
 176                  , node_descendants = m
 177                  , node_size = size m + 1
 178                  })
 179
 180 -- * Extractors
 181
 182 -- | Return the number of non-'Nothing' 'node_content's in the given 'TreeMap'.
 183 size :: Ord k => TreeMap k x -> Int
 184 size = Data.Map.foldr ((+) . node_size) 0
 185
 186 -- | Return the content (if any) associated with the given 'Path'.
 187 find :: Ord k => Path k -> TreeMap k x -> Maybe x
 188 find (k:|[]) m = maybe Nothing node_content $ Data.Map.lookup k m
 189 find (k:|k':ks) m =
 190         maybe Nothing (find (path k' ks) . node_descendants) $
 191         Data.Map.lookup k m
 192
 193 -- | Return the given accumulator folded by the given function
 194 -- applied on non-'Nothing' 'node_content's
 195 -- from left to right through the given 'TreeMap'.
 196 foldlWithKey :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 197 foldlWithKey =
 198         foldp []
 199         where
 200                 foldp :: Ord k
 201                  => [k] -> (a -> Path k -> x -> a)
 202                  -> a -> TreeMap k x -> a
 203                 foldp p fct =
 204                         Data.Map.foldlWithKey
 205                          (\acc k Node{node_content, node_descendants} ->
 206                                 let p' = path k p in
 207                                 let acc' = maybe acc (fct acc (rev p')) node_content in
 208                                 foldp (k:p) fct acc' node_descendants)
 209
 210 -- | Return the given accumulator folded by the given function
 211 -- applied on non-'Nothing' 'node_content's
 212 -- from right to left through the given 'TreeMap'.
 213 foldrWithKey :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 214 foldrWithKey =
 215         foldp []
 216         where
 217                 foldp :: Ord k
 218                  => [k] -> (Path k -> x -> a -> a)
 219                  -> a -> TreeMap k x -> a
 220                 foldp p fct =
 221                         Data.Map.foldrWithKey
 222                          (\k Node{node_content, node_descendants} acc ->
 223                                 let p' = path k p in
 224                                 let acc' = foldp (k:p) fct acc node_descendants in
 225                                 maybe acc' (\x -> fct (rev p') x acc') node_content)
 226
 227 -- | Return a 'Data.Map.Map' associating each 'Path'
 228 -- leading to a non-'Nothing' 'node_content' in the given 'TreeMap',
 229 -- with its content mapped by the given function.
 230 flatten :: Ord k => (x -> y) -> TreeMap k x -> Data.Map.Map (Path k) y
 231 flatten =
 232         flat_map []
 233         where
 234                 flat_map :: Ord k
 235                  => [k] -> (x -> y)
 236                  -> TreeMap k x
 237                  -> Data.Map.Map (Path k) y
 238                 flat_map p f m =
 239                         Data.Map.unions $
 240                         (
 241                         Data.Map.mapKeysMonotonic (rev . flip path p) $
 242                         Data.Map.mapMaybe (\Node{node_content=x} -> f <$> x) m
 243                         ) :
 244                         Data.Map.foldrWithKey
 245                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 246                          [] m