Data/TreeMap/Strict.hs

   1 {-# LANGUAGE DeriveDataTypeable #-}
   2 {-# LANGUAGE NamedFieldPuns #-}
   3 {-# LANGUAGE StandaloneDeriving #-}
   4 {-# LANGUAGE RecordWildCards #-}
   5 {-# OPTIONS_GHC -fno-warn-orphans #-}
   6
   7 -- | This module implements a strict 'TreeMap',
   8 -- which is like a 'Map'
   9 -- but whose key is now a 'NonNull' list of 'Map' keys (a 'Path')
  10 -- enabling the possibility to gather mapped values
  11 -- by 'Path' prefixes (inside a 'Node').
  12 module Data.TreeMap.Strict where
  13
  14 import Control.Applicative (Applicative(..), Alternative((<|>)))
  15 import Control.DeepSeq (NFData(..))
  16 import Control.Monad (Monad(..))
  17 import Data.Bool
  18 import Data.Data (Data)
  19 import Data.Eq (Eq(..))
  20 import Data.Foldable (Foldable, foldMap)
  21 import Data.Function (($), (.), const, flip, id)
  22 import Data.Functor (Functor(..), (<$>))
  23 import Data.Map.Strict (Map)
  24 import Data.Maybe (Maybe(..), maybe)
  25 import Data.Monoid (Monoid(..))
  26 import Data.NonNull (NonNull, nuncons, toNullable)
  27 import Data.Ord (Ord(..))
  28 import Data.Semigroup (Semigroup(..))
  29 import Data.Sequences (reverse)
  30 import Data.Traversable (Traversable(..))
  31 import Data.Typeable (Typeable)
  32 import Prelude (Int, Num(..), seq)
  33 import Text.Show (Show(..))
  34 import qualified Control.Applicative as App
  35 import qualified Data.List as List
  36 import qualified Data.Map.Strict as Map
  37 import qualified Data.NonNull as NonNull
  38 import qualified Data.Strict.Maybe as Strict
  39
  40 -- @Data.Strict@ orphan instances
  41 deriving instance Data x => Data (Strict.Maybe x)
  42 deriving instance Typeable Strict.Maybe
  43 instance Semigroup x => Semigroup (Strict.Maybe x) where
  44         Strict.Just x <> Strict.Just y = Strict.Just (x <> y)
  45         x <> Strict.Nothing = x
  46         Strict.Nothing <> y = y
  47 instance Semigroup x => Monoid (Strict.Maybe x) where
  48         mempty  = Strict.Nothing
  49         mappend = (<>)
  50 instance NFData x => NFData (Strict.Maybe x) where
  51         rnf Strict.Nothing  = ()
  52         rnf (Strict.Just x) = rnf x
  53 instance Applicative Strict.Maybe where
  54         pure = Strict.Just
  55         Strict.Just f <*> Strict.Just x = Strict.Just (f x)
  56         _ <*> _ = Strict.Nothing
  57 instance Alternative Strict.Maybe where
  58         empty = Strict.Nothing
  59         x <|> y = if Strict.isJust x then x else y
  60
  61 -- * Type 'TreeMap'
  62 newtype TreeMap k x
  63  =      TreeMap (Map k (Node k x))
  64  deriving (Data, Eq, Ord, Show, Typeable)
  65
  66 instance (Ord k, Semigroup v) => Semigroup (TreeMap k v) where
  67         (<>) = union (<>)
  68 instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
  69         mempty = empty
  70         mappend = union mappend
  71         -- mconcat = List.foldr mappend mempty
  72 instance Ord k => Functor (TreeMap k) where
  73         fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
  74 instance Ord k => Foldable (TreeMap k) where
  75         foldMap f (TreeMap m) = foldMap (foldMap f) m
  76 instance Ord k => Traversable (TreeMap k) where
  77         traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m
  78 instance (Ord k, NFData k, NFData x) => NFData (TreeMap k x) where
  79         rnf (TreeMap m) = rnf m
  80
  81 -- * Type 'Path'
  82 -- | A 'Path' is a non-empty list of 'Map' keys.
  83 type Path k = NonNull [k]
  84
  85 -- | 'Path' constructor.
  86 path :: k -> [k] -> Path k
  87 path = NonNull.ncons
  88
  89 -- | Convenient alias.
  90 (<|) :: k -> [k] -> Path k
  91 (<|) = path
  92
  93 -- * Type 'Node'
  94 data Node k x
  95  =   Node
  96  {   node_size        :: !Int -- ^ The number of non-'Strict.Nothing' 'node_value's reachable from this 'Node'.
  97  ,   node_value       :: !(Strict.Maybe x) -- ^ Some value, or 'Strict.Nothing' if this 'Node' is intermediary.
  98  ,   node_descendants :: !(TreeMap k x) -- ^ Descendants 'Node's.
  99  } deriving (Data, Eq, Ord, Show, Typeable)
 100
 101 instance (Ord k, Semigroup v) => Semigroup (Node k v) where
 102         (<>)
 103          Node{node_value=x0, node_descendants=m0}
 104          Node{node_value=x1, node_descendants=m1} =
 105                 node (x0 <> x1) (union const m0 m1)
 106 instance (Ord k, Semigroup v) => Monoid (Node k v) where
 107         mempty = node Strict.Nothing (TreeMap mempty)
 108         mappend = (<>)
 109         -- mconcat = List.foldr mappend mempty
 110 instance Ord k => Functor (Node k) where
 111         fmap f Node{node_value=x, node_descendants=m, node_size} =
 112                 Node
 113                  { node_value = fmap f x
 114                  , node_descendants = map f m
 115                  , node_size
 116                  }
 117 instance Ord k => Foldable (Node k) where
 118         foldMap f Node{node_value=Strict.Nothing, node_descendants=TreeMap m} =
 119                 foldMap (foldMap f) m
 120         foldMap f Node{node_value=Strict.Just x, node_descendants=TreeMap m} =
 121                 f x `mappend` foldMap (foldMap f) m
 122 instance Ord k => Traversable (Node k) where
 123         traverse f Node{node_value=Strict.Nothing, node_descendants=TreeMap m, node_size} =
 124                 Node node_size <$> pure Strict.Nothing <*> (TreeMap <$> traverse (traverse f) m)
 125         traverse f Node{node_value=Strict.Just x, node_descendants=TreeMap m, node_size} =
 126                 Node node_size <$> (Strict.Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)
 127 instance (Ord k, NFData k, NFData x) => NFData (Node k x) where
 128         rnf (Node s v d) = rnf s `seq` rnf v `seq` rnf d
 129
 130 node :: Strict.Maybe x -> TreeMap k x -> Node k x
 131 node node_value node_descendants =
 132         Node
 133          { node_value
 134          , node_size =
 135                 size node_descendants +
 136                 Strict.maybe 0 (const 1) node_value
 137          , node_descendants
 138          }
 139
 140 nodeEmpty :: Node k x
 141 nodeEmpty = node Strict.Nothing empty
 142
 143 nodeLookup :: Ord k => [k] -> Node k x -> Strict.Maybe (Node k x)
 144 nodeLookup [] n = Strict.Just n
 145 nodeLookup (k:ks) Node{node_descendants=TreeMap m} =
 146         maybe Strict.Nothing (nodeLookup ks) $
 147         Map.lookup k m
 148
 149 -- * Construct
 150
 151 -- | Return the empty 'TreeMap'.
 152 empty :: TreeMap k x
 153 empty = TreeMap Map.empty
 154
 155 -- | Return a 'TreeMap' only mapping the given 'Path' to the given value.
 156 singleton :: Ord k => Path k -> x -> TreeMap k x
 157 singleton ks x = insert const ks x empty
 158
 159 -- | Return a 'Node' only containing the given value.
 160 leaf :: x -> Node k x
 161 leaf x = node (Strict.Just x) empty
 162
 163 -- | Return the given 'TreeMap' associating the given 'Path' with the given value,
 164 -- merging values if the given 'TreeMap' already associates the given 'Path'
 165 -- with a non-'Strict.Nothing' 'node_value'.
 166 insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
 167 insert merge p x (TreeMap m) =
 168         TreeMap $
 169         case nuncons p of
 170          (k, Nothing) ->
 171                 Map.insertWith (\_ Node{..} -> node
 172                          (Strict.maybe (Strict.Just x) (Strict.Just . merge x) node_value)
 173                          node_descendants)
 174                  k (leaf x) m
 175          (k, Just p') ->
 176                 Map.insertWith (\_ Node{..} -> node node_value $
 177                         insert merge p' x node_descendants)
 178                  k (node Strict.Nothing (insert merge p' x empty)) m
 179
 180 -- | Return a 'TreeMap' from a list of 'Path'/value pairs,
 181 -- with a combining function called on the leftest and rightest values
 182 -- when their 'Path's are identical.
 183 fromList :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
 184 fromList merge = List.foldl' (\acc (p,x) -> insert merge p x acc) empty
 185
 186 -- | Return a 'TreeMap' from a 'Map' mapping 'Path' to value.
 187 fromMap :: Ord k => Map (Path k) x -> TreeMap k x
 188 fromMap = go . Map.toList
 189         where
 190         go :: Ord k => [(Path k,x)] -> TreeMap k x
 191         go m =
 192                 TreeMap $ Map.fromAscListWith
 193                  (\Node{node_value=vn, node_descendants=mn}
 194                    Node{node_value=vo, node_descendants=mo} ->
 195                         node (vn <|> vo) $ union const mn mo) $
 196                 (<$> m) $ \(p,x) ->
 197                         let (p0,mps) = nuncons p in
 198                         case mps of
 199                          Nothing -> (p0,node (Strict.Just x) empty)
 200                          Just ps -> (p0,node Strict.Nothing $ go [(ps,x)])
 201 -- fromMap = Map.foldlWithKey (\acc p x -> insert const p x acc) empty
 202
 203 -- * Size
 204
 205 -- | Return the 'Map' in the given 'TreeMap'.
 206 nodes :: TreeMap k x -> Map k (Node k x)
 207 nodes (TreeMap m) = m
 208
 209 -- | Return 'True' iif. the given 'TreeMap' is of 'size' @0@.
 210 null :: TreeMap k x -> Bool
 211 null m = size m == 0
 212
 213 -- | Return the number of non-'Strict.Nothing' 'node_value's in the given 'TreeMap'.
 214 --
 215 --   * Complexity: O(r) where r is the size of the root 'Map'.
 216 size :: TreeMap k x -> Int
 217 size = Map.foldr ((+) . node_size) 0 . nodes
 218
 219 -- * Find
 220
 221 -- | Return the value (if any) associated with the given 'Path'.
 222 lookup :: Ord k => Path k -> TreeMap k x -> Strict.Maybe x
 223 lookup p (TreeMap m) =
 224         maybe Strict.Nothing nod_val $ Map.lookup k m
 225         where
 226         (k, mp') = nuncons p
 227         nod_val =
 228                 case mp' of
 229                  Nothing -> node_value
 230                  Just p' -> lookup p' . node_descendants
 231
 232 -- | Return the values (if any) associated with the prefixes of the given 'Path' (included).
 233 lookupAlong :: Ord k => Path k -> TreeMap k x -> [x]
 234 lookupAlong p (TreeMap tm) =
 235         go (toNullable p) tm
 236         where
 237                 go :: Ord k => [k] -> Map k (Node k x) -> [x]
 238                 go [] _m = []
 239                 go (k:ks) m =
 240                         case Map.lookup k m of
 241                          Nothing -> []
 242                          Just nod ->
 243                                 Strict.maybe id (:) (node_value nod) $
 244                                 go ks $ nodes (node_descendants nod)
 245
 246 -- | Return the 'Node' (if any) associated with the given 'Path'.
 247 lookupNode :: Ord k => Path k -> TreeMap k x -> Maybe (Node k x)
 248 lookupNode p (TreeMap m) =
 249         case nuncons p of
 250          (k, Nothing) -> Map.lookup k m
 251          (k, Just p') -> Map.lookup k m >>= lookupNode p' . node_descendants
 252
 253 -- * Union
 254
 255 -- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
 256 -- merging values (in respective order) when a 'Path' leads
 257 -- to a non-'Strict.Nothing' 'node_value' in both given 'TreeMap's.
 258 union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
 259 union merge (TreeMap tm0) (TreeMap tm1) =
 260         TreeMap $
 261         Map.unionWith
 262          (\Node{node_value=x0, node_descendants=m0}
 263            Node{node_value=x1, node_descendants=m1} ->
 264                 node (Strict.maybe x1 (\x0' -> Strict.maybe (Strict.Just x0') (Strict.Just . merge x0') x1) x0)
 265                  (union merge m0 m1))
 266          tm0 tm1
 267
 268 -- | Return the 'union' of the given 'TreeMap's.
 269 --
 270 -- NOTE: use |List.foldl'| to reduce demand on the control-stack.
 271 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 272 unions merge = List.foldl' (union merge) empty
 273
 274 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 275 -- foldl' f = go
 276 --      where
 277 --              go z []     = z
 278 --              go z (x:xs) = z `seq` go (f z x) xs
 279
 280 -- * Map
 281
 282 -- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
 283 -- mapped by the given function.
 284 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 285 map f =
 286         TreeMap .
 287         Map.map
 288          (\n@Node{node_value=x, node_descendants=m} ->
 289                 n{ node_value       = fmap f x
 290                  , node_descendants = map f m
 291                  }) .
 292         nodes
 293
 294 -- | Return the given 'TreeMap' with each 'Path' section
 295 -- and each non-'Strict.Nothing' 'node_value'
 296 -- mapped by the given functions.
 297 --
 298 -- WARNING: the function mapping 'Path' sections must be monotonic,
 299 -- like in 'Map.mapKeysMonotonic'.
 300 mapMonotonic :: (Ord k, Ord l) => (k -> l) -> (x -> y) -> TreeMap k x -> TreeMap l y
 301 mapMonotonic fk fx =
 302         TreeMap .
 303         Map.mapKeysMonotonic fk .
 304         Map.map
 305          (\n@Node{node_value=x, node_descendants=m} ->
 306                 n{ node_value       = fmap fx x
 307                  , node_descendants = mapMonotonic fk fx m
 308                  }) .
 309         nodes
 310
 311 -- | Return the given 'TreeMap' with each 'node_value'
 312 -- mapped by the given function supplied with
 313 -- the already mapped 'node_descendants' of the current 'Node'.
 314 mapByDepthFirst :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
 315 mapByDepthFirst f =
 316         TreeMap .
 317         Map.map
 318          (\Node{node_value, node_descendants} ->
 319                 let m = mapByDepthFirst f node_descendants in
 320                 node (Strict.Just $ f m node_value) m) .
 321         nodes
 322
 323 -- * Alter
 324
 325 alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
 326 alterl_path fct =
 327         go fct . toNullable
 328         where
 329                 go :: Ord k
 330                  => (Strict.Maybe x -> Strict.Maybe x) -> [k]
 331                  -> TreeMap k x -> TreeMap k x
 332                 go _f [] m = m
 333                 go f (k:p) (TreeMap m) =
 334                         TreeMap $
 335                         Map.alter
 336                          (\c ->
 337                                 let (cv, cm) =
 338                                         case c of
 339                                          Just Node{node_value=v, node_descendants=d} -> (v, d)
 340                                          Nothing -> (Strict.Nothing, empty) in
 341                                 let fx = f cv in
 342                                 let gm = go f p cm in
 343                                 case (fx, size gm) of
 344                                  (Strict.Nothing, 0) -> Nothing
 345                                  (_, s) -> Just
 346                                         Node
 347                                          { node_value = fx
 348                                          , node_descendants = gm
 349                                          , node_size = s + 1
 350                                          }
 351                          ) k m
 352
 353 -- * Fold
 354
 355 -- | Return the given accumulator folded by the given function
 356 -- applied on non-'Strict.Nothing' 'node_value's
 357 -- from left to right through the given 'TreeMap'.
 358 foldlWithPath :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 359 foldlWithPath =
 360         foldp []
 361         where
 362                 foldp :: Ord k
 363                  => [k] -> (a -> Path k -> x -> a)
 364                  -> a -> TreeMap k x -> a
 365                 foldp p fct a (TreeMap m) =
 366                         Map.foldlWithKey
 367                          (\acc k Node{..} ->
 368                                 let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
 369                                 foldp (k:p) fct acc' node_descendants) a m
 370
 371 -- | Return the given accumulator folded by the given function
 372 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
 373 -- from left to right through the given 'TreeMap'.
 374 foldlWithPathAndNode :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
 375 foldlWithPathAndNode =
 376         foldp []
 377         where
 378                 foldp :: Ord k
 379                  => [k] -> (a -> Node k x -> Path k -> x -> a)
 380                  -> a -> TreeMap k x -> a
 381                 foldp p fct a (TreeMap m) =
 382                         Map.foldlWithKey
 383                          (\acc k n@Node{..} ->
 384                                 let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
 385                                 foldp (k:p) fct acc' node_descendants) a m
 386
 387 -- | Return the given accumulator folded by the given function
 388 -- applied on non-'Strict.Nothing' 'node_value's
 389 -- from right to left through the given 'TreeMap'.
 390 foldrWithPath :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 391 foldrWithPath =
 392         foldp []
 393         where
 394                 foldp :: Ord k
 395                  => [k] -> (Path k -> x -> a -> a)
 396                  -> a -> TreeMap k x -> a
 397                 foldp p fct a (TreeMap m) =
 398                         Map.foldrWithKey
 399                          (\k Node{..} acc ->
 400                                 let acc' = foldp (k:p) fct acc node_descendants in
 401                                 Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
 402
 403 -- | Return the given accumulator folded by the given function
 404 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
 405 -- from right to left through the given 'TreeMap'.
 406 foldrWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 407 foldrWithPathAndNode =
 408         foldp []
 409         where
 410                 foldp :: Ord k
 411                  => [k] -> (Node k x -> Path k -> x -> a -> a)
 412                  -> a -> TreeMap k x -> a
 413                 foldp p fct a (TreeMap m) =
 414                         Map.foldrWithKey
 415                          (\k n@Node{..} acc ->
 416                                 let acc' = foldp (k:p) fct acc node_descendants in
 417                                 Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
 418
 419 -- | Return the given accumulator folded by the given function
 420 -- applied on non-'Strict.Nothing' 'node_value's
 421 -- from left to right along the given 'Path'.
 422 foldlPath :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 423 foldlPath fct =
 424         go fct [] . toNullable
 425         where
 426                 go :: Ord k
 427                  => (Path k -> x -> a -> a) -> [k] -> [k]
 428                  -> TreeMap k x -> a -> a
 429                 go _f _ [] _t a = a
 430                 go f p (k:n) (TreeMap t) a =
 431                         case Map.lookup k t of
 432                          Nothing -> a
 433                          Just Node{..} ->
 434                                 case node_value of
 435                                  Strict.Nothing -> go f (k:p) n node_descendants a
 436                                  Strict.Just x  -> go f (k:p) n node_descendants (f (reverse $ path k p) x a)
 437
 438 -- | Return the given accumulator folded by the given function
 439 -- applied on non-'Strict.Nothing' 'node_value's
 440 -- from right to left along the given 'Path'.
 441 foldrPath :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 442 foldrPath fct =
 443         go fct [] . toNullable
 444         where
 445                 go :: Ord k
 446                  => (Path k -> x -> a -> a) -> [k] -> [k]
 447                  -> TreeMap k x -> a -> a
 448                 go _f _ [] _t a = a
 449                 go f p (k:n) (TreeMap t) a =
 450                         case Map.lookup k t of
 451                          Nothing -> a
 452                          Just Node{..} ->
 453                                 case node_value of
 454                                  Strict.Nothing -> go f (k:p) n node_descendants a
 455                                  Strict.Just x  -> f (reverse $ path k p) x $ go f (k:p) n node_descendants a
 456
 457 -- * Flatten
 458
 459 -- | Return a 'Map' associating each 'Path'
 460 -- leading to a non-'Strict.Nothing' 'node_value' in the given 'TreeMap',
 461 -- with its value mapped by the given function.
 462 flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
 463 flatten = flattenWithPath . const
 464
 465 -- | Like 'flatten' but with also the current 'Path' given to the mapping function.
 466 flattenWithPath :: Ord k => (Path k -> x -> y) -> TreeMap k x -> Map (Path k) y
 467 flattenWithPath =
 468         flat_map []
 469         where
 470                 flat_map :: Ord k
 471                  => [k] -> (Path k -> x -> y)
 472                  -> TreeMap k x
 473                  -> Map (Path k) y
 474                 flat_map p f (TreeMap m) =
 475                         Map.unions $
 476                         Map.mapKeysMonotonic (reverse . flip path p) (
 477                         Map.mapMaybeWithKey (\k Node{node_value} ->
 478                                 case node_value of
 479                                  Strict.Nothing -> Nothing
 480                                  Strict.Just x  -> Just $ f (reverse $ path k p) x) m
 481                         ) :
 482                         Map.foldrWithKey
 483                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 484                          [] m
 485
 486 -- * Filter
 487
 488 -- | Return the given 'TreeMap'
 489 --   keeping only its non-'Strict.Nothing' 'node_value's
 490 --   passing the given predicate.
 491 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
 492 filter f =
 493         mapMaybeWithPath
 494          (\_p x -> if f x then Strict.Just x else Strict.Nothing)
 495
 496 -- | Like 'filter' but with also the current 'Path' given to the predicate.
 497 filterWithPath :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 498 filterWithPath f =
 499         mapMaybeWithPath
 500          (\p x -> if f p x then Strict.Just x else Strict.Nothing)
 501
 502 -- | Like 'filterWithPath' but with also the current 'Node' given to the predicate.
 503 filterWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 504 filterWithPathAndNode f =
 505         mapMaybeWithPathAndNode
 506          (\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
 507
 508 -- | Return the given 'TreeMap'
 509 --   mapping its non-'Strict.Nothing' 'node_value's
 510 --   and keeping only the non-'Strict.Nothing' results.
 511 mapMaybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 512 mapMaybe = mapMaybeWithPath . const
 513
 514 -- | Like 'mapMaybe' but with also the current 'Path' given to the predicate.
 515 mapMaybeWithPath :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 516 mapMaybeWithPath = mapMaybeWithPathAndNode . const
 517
 518 -- | Like 'mapMaybeWithPath' but with also the current 'Node' given to the predicate.
 519 mapMaybeWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 520 mapMaybeWithPathAndNode =
 521         go []
 522         where
 523                 go :: Ord k
 524                  => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
 525                  -> TreeMap k x
 526                  -> TreeMap k y
 527                 go p test (TreeMap m) =
 528                         TreeMap $
 529                         Map.mapMaybeWithKey
 530                          (\k nod@Node{node_value=v, node_descendants=ns} ->
 531                                 let node_descendants = go (k:p) test ns in
 532                                 let node_size = size node_descendants in
 533                                 case v of
 534                                  Strict.Just x ->
 535                                         let node_value = test nod (reverse $ path k p) x in
 536                                         case node_value of
 537                                          Strict.Nothing | null node_descendants -> Nothing
 538                                          Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
 539                                          Strict.Just _  -> Just Node{node_value, node_descendants, node_size}
 540                                  _ ->
 541                                         if null node_descendants
 542                                         then Nothing
 543                                         else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
 544                          ) m
 545
 546 -- * Intersection
 547
 548 (\\) :: Ord k => TreeMap k x -> TreeMap k y -> TreeMap k x
 549 (\\) = intersection const
 550
 551 intersection ::
 552  Ord k =>
 553  (Strict.Maybe x -> Strict.Maybe y -> Strict.Maybe z) ->
 554  TreeMap k x -> TreeMap k y -> TreeMap k z
 555 intersection merge (TreeMap x) (TreeMap y) =
 556         TreeMap $
 557         Map.intersectionWith
 558          (\xn yn ->
 559                 node (node_value xn `merge` node_value yn) $
 560                 intersection merge
 561                  (node_descendants xn)
 562                  (node_descendants yn))
 563          x y