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 'empty'.
 210 null :: TreeMap k x -> Bool
 211 null (TreeMap m) = Map.null m
 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
 269
 270 -- | Return the 'union' of the given 'TreeMap's.
 271 --
 272 -- NOTE: use |List.foldl'| to reduce demand on the control-stack.
 273 unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
 274 unions merge = List.foldl' (union merge) empty
 275
 276 -- foldl' :: (a -> b -> a) -> a -> [b] -> a
 277 -- foldl' f = go
 278 --      where
 279 --              go z []     = z
 280 --              go z (x:xs) = z `seq` go (f z x) xs
 281
 282 -- * Map
 283
 284 -- | Return the given 'TreeMap' with each non-'Strict.Nothing' 'node_value'
 285 -- mapped by the given function.
 286 map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
 287 map f =
 288         TreeMap .
 289         Map.map
 290          (\n@Node{node_value=x, node_descendants=m} ->
 291                 n{ node_value       = fmap f x
 292                  , node_descendants = map f m
 293                  }) .
 294         nodes
 295
 296 -- | Return the given 'TreeMap' with each 'Path' section
 297 -- and each non-'Strict.Nothing' 'node_value'
 298 -- mapped by the given functions.
 299 --
 300 -- WARNING: the function mapping 'Path' sections must be monotonic,
 301 -- like in 'Map.mapKeysMonotonic'.
 302 mapMonotonic :: (Ord k, Ord l) => (k -> l) -> (x -> y) -> TreeMap k x -> TreeMap l y
 303 mapMonotonic fk fx =
 304         TreeMap .
 305         Map.mapKeysMonotonic fk .
 306         Map.map
 307          (\n@Node{node_value=x, node_descendants=m} ->
 308                 n{ node_value       = fmap fx x
 309                  , node_descendants = mapMonotonic fk fx m
 310                  }) .
 311         nodes
 312
 313 -- | Return the given 'TreeMap' with each 'node_value'
 314 -- mapped by the given function supplied with
 315 -- the already mapped 'node_descendants' of the current 'Node'.
 316 mapByDepthFirst :: Ord k => (TreeMap k y -> Strict.Maybe x -> y) -> TreeMap k x -> TreeMap k y
 317 mapByDepthFirst f =
 318         TreeMap .
 319         Map.map
 320          (\Node{node_value, node_descendants} ->
 321                 let m = mapByDepthFirst f node_descendants in
 322                 node (Strict.Just $ f m node_value) m) .
 323         nodes
 324
 325 -- * Alter
 326
 327 alterl_path :: Ord k => (Strict.Maybe x -> Strict.Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
 328 alterl_path fct =
 329         go fct . toNullable
 330         where
 331                 go :: Ord k
 332                  => (Strict.Maybe x -> Strict.Maybe x) -> [k]
 333                  -> TreeMap k x -> TreeMap k x
 334                 go _f [] m = m
 335                 go f (k:p) (TreeMap m) =
 336                         TreeMap $
 337                         Map.alter
 338                          (\c ->
 339                                 let (cv, cm) =
 340                                         case c of
 341                                          Just Node{node_value=v, node_descendants=d} -> (v, d)
 342                                          Nothing -> (Strict.Nothing, empty) in
 343                                 let fx = f cv in
 344                                 let gm = go f p cm in
 345                                 case (fx, size gm) of
 346                                  (Strict.Nothing, 0) -> Nothing
 347                                  (_, s) -> Just
 348                                         Node
 349                                          { node_value = fx
 350                                          , node_descendants = gm
 351                                          , node_size = s + 1
 352                                          }
 353                          ) k m
 354
 355 -- * Fold
 356
 357 -- | Return the given accumulator folded by the given function
 358 -- applied on non-'Strict.Nothing' 'node_value's
 359 -- from left to right through the given 'TreeMap'.
 360 foldlWithPath :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
 361 foldlWithPath =
 362         foldp []
 363         where
 364                 foldp :: Ord k
 365                  => [k] -> (a -> Path k -> x -> a)
 366                  -> a -> TreeMap k x -> a
 367                 foldp p fct a (TreeMap m) =
 368                         Map.foldlWithKey
 369                          (\acc k Node{..} ->
 370                                 let acc' = Strict.maybe acc (fct acc (reverse $ path k p)) node_value in
 371                                 foldp (k:p) fct acc' node_descendants) a m
 372
 373 -- | Return the given accumulator folded by the given function
 374 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
 375 -- from left to right through the given 'TreeMap'.
 376 foldlWithPathAndNode :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
 377 foldlWithPathAndNode =
 378         foldp []
 379         where
 380                 foldp :: Ord k
 381                  => [k] -> (a -> Node k x -> Path k -> x -> a)
 382                  -> a -> TreeMap k x -> a
 383                 foldp p fct a (TreeMap m) =
 384                         Map.foldlWithKey
 385                          (\acc k n@Node{..} ->
 386                                 let acc' = Strict.maybe acc (fct acc n (reverse $ path k p)) node_value in
 387                                 foldp (k:p) fct acc' node_descendants) a m
 388
 389 -- | Return the given accumulator folded by the given function
 390 -- applied on non-'Strict.Nothing' 'node_value's
 391 -- from right to left through the given 'TreeMap'.
 392 foldrWithPath :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 393 foldrWithPath =
 394         foldp []
 395         where
 396                 foldp :: Ord k
 397                  => [k] -> (Path k -> x -> a -> a)
 398                  -> a -> TreeMap k x -> a
 399                 foldp p fct a (TreeMap m) =
 400                         Map.foldrWithKey
 401                          (\k Node{..} acc ->
 402                                 let acc' = foldp (k:p) fct acc node_descendants in
 403                                 Strict.maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m
 404
 405 -- | Return the given accumulator folded by the given function
 406 -- applied on non-'Strict.Nothing' 'Node's and 'node_value's
 407 -- from right to left through the given 'TreeMap'.
 408 foldrWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
 409 foldrWithPathAndNode =
 410         foldp []
 411         where
 412                 foldp :: Ord k
 413                  => [k] -> (Node k x -> Path k -> x -> a -> a)
 414                  -> a -> TreeMap k x -> a
 415                 foldp p fct a (TreeMap m) =
 416                         Map.foldrWithKey
 417                          (\k n@Node{..} acc ->
 418                                 let acc' = foldp (k:p) fct acc node_descendants in
 419                                 Strict.maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m
 420
 421 -- | Return the given accumulator folded by the given function
 422 -- applied on non-'Strict.Nothing' 'node_value's
 423 -- from left to right along the given 'Path'.
 424 foldlPath :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 425 foldlPath fct =
 426         go fct [] . toNullable
 427         where
 428                 go :: Ord k
 429                  => (Path k -> x -> a -> a) -> [k] -> [k]
 430                  -> TreeMap k x -> a -> a
 431                 go _f _ [] _t a = a
 432                 go f p (k:n) (TreeMap t) a =
 433                         case Map.lookup k t of
 434                          Nothing -> a
 435                          Just Node{..} ->
 436                                 case node_value of
 437                                  Strict.Nothing -> go f (k:p) n node_descendants a
 438                                  Strict.Just x  -> go f (k:p) n node_descendants (f (reverse $ path k p) x a)
 439
 440 -- | Return the given accumulator folded by the given function
 441 -- applied on non-'Strict.Nothing' 'node_value's
 442 -- from right to left along the given 'Path'.
 443 foldrPath :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
 444 foldrPath fct =
 445         go fct [] . toNullable
 446         where
 447                 go :: Ord k
 448                  => (Path k -> x -> a -> a) -> [k] -> [k]
 449                  -> TreeMap k x -> a -> a
 450                 go _f _ [] _t a = a
 451                 go f p (k:n) (TreeMap t) a =
 452                         case Map.lookup k t of
 453                          Nothing -> a
 454                          Just Node{..} ->
 455                                 case node_value of
 456                                  Strict.Nothing -> go f (k:p) n node_descendants a
 457                                  Strict.Just x  -> f (reverse $ path k p) x $ go f (k:p) n node_descendants a
 458
 459 -- * Flatten
 460
 461 -- | Return a 'Map' associating each 'Path'
 462 -- leading to a non-'Strict.Nothing' 'node_value' in the given 'TreeMap',
 463 -- with its value mapped by the given function.
 464 flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
 465 flatten = flattenWithPath . const
 466
 467 -- | Like 'flatten' but with also the current 'Path' given to the mapping function.
 468 flattenWithPath :: Ord k => (Path k -> x -> y) -> TreeMap k x -> Map (Path k) y
 469 flattenWithPath =
 470         flat_map []
 471         where
 472                 flat_map :: Ord k
 473                  => [k] -> (Path k -> x -> y)
 474                  -> TreeMap k x
 475                  -> Map (Path k) y
 476                 flat_map p f (TreeMap m) =
 477                         Map.unions $
 478                         Map.mapKeysMonotonic (reverse . flip path p) (
 479                         Map.mapMaybeWithKey (\k Node{node_value} ->
 480                                 case node_value of
 481                                  Strict.Nothing -> Nothing
 482                                  Strict.Just x  -> Just $ f (reverse $ path k p) x) m
 483                         ) :
 484                         Map.foldrWithKey
 485                          (\k -> (:) . flat_map (k:p) f . node_descendants)
 486                          [] m
 487
 488 -- * Filter
 489
 490 -- | Return the given 'TreeMap'
 491 --   keeping only its non-'Strict.Nothing' 'node_value's
 492 --   passing the given predicate.
 493 filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
 494 filter f =
 495         mapMaybeWithPath
 496          (\_p x -> if f x then Strict.Just x else Strict.Nothing)
 497
 498 -- | Like 'filter' but with also the current 'Path' given to the predicate.
 499 filterWithPath :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 500 filterWithPath f =
 501         mapMaybeWithPath
 502          (\p x -> if f p x then Strict.Just x else Strict.Nothing)
 503
 504 -- | Like 'filterWithPath' but with also the current 'Node' given to the predicate.
 505 filterWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
 506 filterWithPathAndNode f =
 507         mapMaybeWithPathAndNode
 508          (\n p x -> if f n p x then Strict.Just x else Strict.Nothing)
 509
 510 -- | Return the given 'TreeMap'
 511 --   mapping its non-'Strict.Nothing' 'node_value's
 512 --   and keeping only the non-'Strict.Nothing' results.
 513 mapMaybe :: Ord k => (x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 514 mapMaybe = mapMaybeWithPath . const
 515
 516 -- | Like 'mapMaybe' but with also the current 'Path' given to the predicate.
 517 mapMaybeWithPath :: Ord k => (Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 518 mapMaybeWithPath = mapMaybeWithPathAndNode . const
 519
 520 -- | Like 'mapMaybeWithPath' but with also the current 'Node' given to the predicate.
 521 mapMaybeWithPathAndNode :: Ord k => (Node k x -> Path k -> x -> Strict.Maybe y) -> TreeMap k x -> TreeMap k y
 522 mapMaybeWithPathAndNode =
 523         go []
 524         where
 525                 go :: Ord k
 526                  => [k] -> (Node k x -> Path k -> x -> Strict.Maybe y)
 527                  -> TreeMap k x
 528                  -> TreeMap k y
 529                 go p test (TreeMap m) =
 530                         TreeMap $
 531                         Map.mapMaybeWithKey
 532                          (\k nod@Node{node_value=v, node_descendants=ns} ->
 533                                 let node_descendants = go (k:p) test ns in
 534                                 let node_size = size node_descendants in
 535                                 case v of
 536                                  Strict.Just x ->
 537                                         let node_value = test nod (reverse $ path k p) x in
 538                                         case node_value of
 539                                          Strict.Nothing | null node_descendants -> Nothing
 540                                          Strict.Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
 541                                          Strict.Just _  -> Just Node{node_value, node_descendants, node_size}
 542                                  _ ->
 543                                         if null node_descendants
 544                                         then Nothing
 545                                         else Just Node{node_value=Strict.Nothing, node_descendants, node_size}
 546                          ) m