From f81a9e69ff76c4a86e329cd531164534977cee17 Mon Sep 17 00:00:00 2001 From: Tom Smalley Date: Tue, 21 Apr 2020 17:20:01 +0100 Subject: [PATCH] Add merge --- dependent-map.cabal | 3 +- src/Data/Dependent/Map.hs | 72 +----- src/Data/Dependent/Map/Internal.hs | 106 +++++++- src/Data/Dependent/Map/Merge.hs | 403 +++++++++++++++++++++++++++++ 4 files changed, 512 insertions(+), 72 deletions(-) create mode 100644 src/Data/Dependent/Map/Merge.hs diff --git a/dependent-map.cabal b/dependent-map.cabal index dc51238..1f1bc98 100644 --- a/dependent-map.cabal +++ b/dependent-map.cabal @@ -1,5 +1,5 @@ name: dependent-map -version: 0.4.0.0 +version: 0.4.1.0 stability: provisional cabal-version: >= 1.6 @@ -35,6 +35,7 @@ Library ghc-options: -fwarn-unused-imports -fwarn-unused-binds exposed-modules: Data.Dependent.Map, Data.Dependent.Map.Lens, + Data.Dependent.Map.Merge, Data.Dependent.Map.Internal other-modules: Data.Dependent.Map.PtrEquality build-depends: base >= 4.9 && < 5, diff --git a/src/Data/Dependent/Map.hs b/src/Data/Dependent/Map.hs index fd0e17d..09ae3e2 100644 --- a/src/Data/Dependent/Map.hs +++ b/src/Data/Dependent/Map.hs @@ -63,6 +63,9 @@ module Data.Dependent.Map , intersection , intersectionWithKey + -- ** General combining functions + -- | See "Data.Dependent.Map.Merge" + -- * Traversal -- ** Map , map @@ -72,6 +75,7 @@ module Data.Dependent.Map , traverseWithKey_ , forWithKey_ , traverseWithKey + , traverseMaybeWithKey , forWithKey , mapAccumLWithKey , mapAccumRWithKey @@ -659,7 +663,7 @@ difference t1 Tip = t1 difference t1 (Bin _ k2 _x2 l2 r2) = case split k2 t1 of (l1, r1) | size t1 == size l1l2 + size r1r2 -> t1 - | otherwise -> merge l1l2 r1r2 + | otherwise -> link2 l1l2 r1r2 where !l1l2 = l1 `difference` l2 !r1r2 = r1 `difference` r2 @@ -675,7 +679,7 @@ differenceWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of (l2, mx2, r2) -> case mx2 of Nothing -> combine k1 x1 l1l2 r1r2 Just x2 -> case f k1 x1 x2 of - Nothing -> merge l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 Just x1x2 -> combine k1 x1x2 l1l2 r1r2 where !l1l2 = differenceWithKey f l1 l2 !r1r2 = differenceWithKey f r1 r2 @@ -698,7 +702,7 @@ intersection t1@(Bin s1 k1 x1 l1 r1) t2 = then if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 then t1 else combine k1 x1 l1l2 r1r2 - else merge l1l2 r1r2 + else link2 l1l2 r1r2 -- | /O(m * log (n\/m + 1), m <= n/. Intersection with a combining function. intersectionWithKey :: GCompare k => (forall v. k v -> f v -> g v -> h v) -> DMap k f -> DMap k g -> DMap k h @@ -709,7 +713,7 @@ intersectionWithKey f (Bin s1 k1 x1 l1 r1) t2 = !l1l2 = intersectionWithKey f l1 l2 !r1r2 = intersectionWithKey f r1 r2 in case found of - Nothing -> merge l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 Just x2 -> combine k1 (f k1 x1 x2) l1l2 r1r2 {-------------------------------------------------------------------- @@ -766,19 +770,6 @@ isProperSubmapOfBy f t1 t2 Filter and partition --------------------------------------------------------------------} --- | /O(n)/. Filter all keys\/values that satisfy the predicate. -filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> DMap k f -filterWithKey p = go - where - go Tip = Tip - go t@(Bin _ kx x l r) - | p kx x = if l' `ptrEq` l && r' `ptrEq` r - then t - else combine kx x l' r' - | otherwise = merge l' r' - where !l' = go l - !r' = go r - -- | /O(n)/. Partition the map according to a predicate. The first -- map contains all elements that satisfy the predicate, the second all -- elements that fail the predicate. See also 'split'. @@ -788,8 +779,8 @@ partitionWithKey p0 m0 = toPair (go p0 m0) go :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> (DMap k f :*: DMap k f) go _ Tip = (Tip :*: Tip) go p (Bin _ kx x l r) - | p kx x = (combine kx x l1 r1 :*: merge l2 r2) - | otherwise = (merge l1 r1 :*: combine kx x l2 r2) + | p kx x = (combine kx x l1 r1 :*: link2 l2 r2) + | otherwise = (link2 l1 r1 :*: combine kx x l2 r2) where (l1 :*: l2) = go p l (r1 :*: r2) = go p r @@ -798,15 +789,6 @@ partitionWithKey p0 m0 = toPair (go p0 m0) mapMaybe :: GCompare k => (forall v. f v -> Maybe (g v)) -> DMap k f -> DMap k g mapMaybe f = mapMaybeWithKey (const f) --- | /O(n)/. Map keys\/values and collect the 'Just' results. -mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> DMap k f -> DMap k g -mapMaybeWithKey f = go - where - go Tip = Tip - go (Bin _ kx x l r) = case f kx x of - Just y -> combine kx y (go l) (go r) - Nothing -> merge (go l) (go r) - -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. mapEitherWithKey :: GCompare k => (forall v. k v -> f v -> Either (g v) (h v)) -> DMap k f -> (DMap k g, DMap k h) @@ -817,8 +799,8 @@ mapEitherWithKey f0 = toPair . go f0 -> DMap k f -> (DMap k g :*: DMap k h) go _ Tip = (Tip :*: Tip) go f (Bin _ kx x l r) = case f kx x of - Left y -> (combine kx y l1 r1 :*: merge l2 r2) - Right z -> (merge l1 r1 :*: combine kx z l2 r2) + Left y -> (combine kx y l1 r1 :*: link2 l2 r2) + Right z -> (link2 l1 r1 :*: combine kx z l2 r2) where (l1,l2) = mapEitherWithKey f l (r1,r2) = mapEitherWithKey f r @@ -840,13 +822,6 @@ map f = go ffor :: DMap k f -> (forall v. f v -> g v) -> DMap k g ffor m f = map f m --- | /O(n)/. Map a function over all values in the map. -mapWithKey :: (forall v. k v -> f v -> g v) -> DMap k f -> DMap k g -mapWithKey f = go - where - go Tip = Tip - go (Bin sx kx x l r) = Bin sx kx (f kx x) (go l) (go r) - -- | /O(n)/. -- @'fforWithKey' == 'flip' 'mapWithKey'@ except we cannot actually use -- 'flip' because of the lack of impredicative types. @@ -870,17 +845,6 @@ traverseWithKey_ f = go forWithKey_ :: Applicative t => DMap k f -> (forall v. k v -> f v -> t ()) -> t () forWithKey_ m f = traverseWithKey_ f m --- | /O(n)/. --- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ --- That is, behaves exactly like a regular 'traverse' except that the traversing --- function also has access to the key associated with a value. -traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> DMap k f -> t (DMap k g) -traverseWithKey f = go - where - go Tip = pure Tip - go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v - go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r - -- | /O(n)/. -- @'forWithKey' == 'flip' 'traverseWithKey'@ except we cannot actually use -- 'flip' because of the lack of impredicative types. @@ -1112,18 +1076,6 @@ split k = toPair . go GEQ -> (l :*: r) {-# INLINABLE split #-} --- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just --- like 'split' but also returns @'lookup' k map@. -splitLookup :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (f v), DMap k f) -splitLookup k = toTriple . go - where - go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f) - go Tip = Triple' Tip Nothing Tip - go (Bin _ kx x l r) = case gcompare k kx of - GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r) - GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt - GEQ -> Triple' l (Just x) r - -- | /O(log n)/. The expression (@'splitMember' k map@) splits a map just -- like 'split' but also returns @'member' k map@. splitMember :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Bool, DMap k f) diff --git a/src/Data/Dependent/Map/Internal.hs b/src/Data/Dependent/Map/Internal.hs index 3bb9487..2adbe48 100644 --- a/src/Data/Dependent/Map/Internal.hs +++ b/src/Data/Dependent/Map/Internal.hs @@ -1,17 +1,21 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DeriveDataTypeable #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE PolyKinds #-} -{-# LANGUAGE Safe #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-} module Data.Dependent.Map.Internal where +import Control.Applicative (liftA3) import Data.Dependent.Sum (DSum((:=>))) import Data.GADT.Compare (GCompare, GOrdering(..), gcompare) import Data.Some (Some, mkSome, withSome) import Data.Typeable (Typeable) +import Data.Dependent.Map.PtrEquality (ptrEq) + -- |Dependent maps: 'k' is a GADT-like thing with a facility for -- rediscovering its type parameter, elements of which function as identifiers -- tagged with the type of the thing they identify. Real GADTs are one @@ -34,6 +38,87 @@ data DMap k f where -> DMap k f deriving Typeable +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> DMap k f +filterWithKey p = go + where + go Tip = Tip + go t@(Bin _ kx x l r) + | p kx x = if l' `ptrEq` l && r' `ptrEq` r + then t + else combine kx x l' r' + | otherwise = link2 l' r' + where !l' = go l + !r' = go r + +-- | /O(n)/. Filter keys and values using an 'Applicative' +-- predicate. +filterWithKeyA :: (GCompare k, Applicative t) => (forall v. k v -> f v -> t Bool) -> DMap k f -> t (DMap k f) +filterWithKeyA _ Tip = pure Tip +filterWithKeyA p t@(Bin _ kx x l r) = + liftA3 combine' (p kx x) (filterWithKeyA p l) (filterWithKeyA p r) + where + combine' True pl pr + | pl `ptrEq` l && pr `ptrEq` r = t + | otherwise = combine kx x pl pr + combine' False pl pr = link2 pl pr + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> DMap k f -> DMap k g +mapMaybeWithKey f = go + where + go Tip = Tip + go (Bin _ kx x l r) = case f kx x of + Just y -> combine kx y (go l) (go r) + Nothing -> link2 (go l) (go r) + +-- | /O(n)/. Map a function over all values in the map. +mapWithKey :: (forall v. k v -> f v -> g v) -> DMap k f -> DMap k g +mapWithKey f = go + where + go Tip = Tip + go (Bin sx kx x l r) = Bin sx kx (f kx x) (go l) (go r) + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- That is, behaves exactly like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value. +traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> DMap k f -> t (DMap k g) +traverseWithKey f = go + where + go Tip = pure Tip + go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v + go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r + +-- | /O(n)/. Traverse keys\/values and collect the 'Just' results. +-- +-- @since UNRELEASED +traverseMaybeWithKey + :: (GCompare k, Applicative t) + => (forall v. k v -> f v -> t (Maybe (g v))) -> DMap k f -> t (DMap k g) +traverseMaybeWithKey f = go + where + go Tip = pure Tip + go (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x + go (Bin _ kx x l r) = liftA3 combine' (go l) (f kx x) (go r) + where + combine' !l' mx !r' = case mx of + Nothing -> link2 l' r' + Just x' -> combine kx x' l' r' + + +-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@. +splitLookup :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (f v), DMap k f) +splitLookup k = toTriple . go + where + go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f) + go Tip = Triple' Tip Nothing Tip + go (Bin _ kx x l r) = case gcompare k kx of + GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r) + GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt + GEQ -> Triple' l (Just x) r + {-------------------------------------------------------------------- Construction --------------------------------------------------------------------} @@ -112,10 +197,10 @@ lookupAssoc sk = withSome sk $ \k -> are valid: [glue l r] Glues [l] and [r] together. Assumes that [l] and [r] are already balanced with respect to each other. - [merge l r] Merges two trees and restores balance. + [link2 l r] Merges two trees and restores balance. Note: in contrast to Adam's paper, we use (<=) comparisons instead - of (<) comparisons in [combine], [merge] and [balance]. + of (<) comparisons in [combine], [link2] and [balance]. Quickcheck (on [difference]) showed that this was necessary in order to maintain the invariants. It is quite unsatisfactory that I haven't been able to find out why this is actually the case! Fortunately, it @@ -133,7 +218,6 @@ combine kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) | delta*sizeR <= sizeL = balance ky y ly (combine kx x ry r) | otherwise = bin kx x l r - -- insertMin and insertMax don't perform potentially expensive comparisons. insertMax,insertMin :: k v -> f v -> DMap k f -> DMap k f insertMax kx x t @@ -149,14 +233,14 @@ insertMin kx x t -> balance ky y (insertMin kx x l) r {-------------------------------------------------------------------- - [merge l r]: merges two trees. + [link2 l r]: merges two trees. --------------------------------------------------------------------} -merge :: DMap k f -> DMap k f -> DMap k f -merge Tip r = r -merge l Tip = l -merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) - | delta*sizeL <= sizeR = balance ky y (merge l ly) ry - | delta*sizeR <= sizeL = balance kx x lx (merge rx r) +link2 :: DMap k f -> DMap k f -> DMap k f +link2 Tip r = r +link2 l Tip = l +link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) + | delta*sizeL <= sizeR = balance ky y (link2 l ly) ry + | delta*sizeR <= sizeL = balance kx x lx (link2 rx r) | otherwise = glue l r {-------------------------------------------------------------------- diff --git a/src/Data/Dependent/Map/Merge.hs b/src/Data/Dependent/Map/Merge.hs new file mode 100644 index 0000000..b13d923 --- /dev/null +++ b/src/Data/Dependent/Map/Merge.hs @@ -0,0 +1,403 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +module Data.Dependent.Map.Merge + ( + -- ** Simple merge tactic types + SimpleWhenMissing + , SimpleWhenMatched + + -- ** General combining function + , merge + + -- *** @WhenMatched@ tactics + , zipWithMaybeMatched + , zipWithMatched + + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative merge tactic types + , WhenMissing + , WhenMatched + + -- ** Applicative general combining function + , mergeA + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + + -- *** Covariant maps for tactics + , mapWhenMissing + , mapWhenMatched + + -- *** Contravariant maps for tactics + , lmapWhenMissing + , contramapFirstWhenMatched + , contramapSecondWhenMatched + + -- *** Miscellaneous tactic functions + , runWhenMatched + , runWhenMissing + ) where + +import Data.GADT.Compare (GCompare) +import Control.Applicative (liftA3) +import Data.Functor.Identity +import Data.Dependent.Map.Internal + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @ WhenMissing f k g h @ is an abstract representation +-- of a function of type @ forall a. k a -> g a -> f (Maybe (h a)) @. +-- +-- @since UNRELEASED +data WhenMissing f k x y = WhenMissing + { missingSubtree :: DMap k x -> f (DMap k y) + , missingKey :: forall a. k a -> x a -> f (Maybe (y a)) + } + +-- | Map covariantly over a @'WhenMissing' f k x@. +-- +-- @since UNRELEASED +mapWhenMissing + :: (Applicative f, Monad f) + => (forall a. g a -> h a) + -> WhenMissing f k x g -> WhenMissing f k x h +mapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! mapWithKey (\_ -> f) m' + , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) } +{-# INLINE mapWhenMissing #-} + +-- | Map contravariantly over a @'WhenMissing' f k _ x@. +-- +-- @since UNRELEASED +lmapWhenMissing :: (forall a. h a -> g a) -> WhenMissing f k g x -> WhenMissing f k h x +lmapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t (mapWithKey (const f) m) + , missingKey = \k x -> missingKey t k (f x) } +{-# INLINE lmapWhenMissing #-} + +-- | Map contravariantly over a @'WhenMatched' f k _ y z@. +-- +-- @since UNRELEASED +contramapFirstWhenMatched + :: (forall a. h a -> g a) + -> WhenMatched f k g y z + -> WhenMatched f k h y z +contramapFirstWhenMatched f t = WhenMatched $ + \k x y -> runWhenMatched t k (f x) y +{-# INLINE contramapFirstWhenMatched #-} + +-- | Map contravariantly over a @'WhenMatched' f k x _ z@. +-- +-- @since UNRELEASED +contramapSecondWhenMatched + :: (forall a. h a -> g a) + -> WhenMatched f k x g z + -> WhenMatched f k x h z +contramapSecondWhenMatched f t = WhenMatched $ + \k x y -> runWhenMatched t k x (f y) +{-# INLINE contramapSecondWhenMatched #-} + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge'. +-- +-- A tactic of type @ SimpleWhenMissing k x z @ is an abstract representation +-- of a function of type @ forall a. k a -> x a -> Maybe (z a) @. +-- +-- @since UNRELEASED +type SimpleWhenMissing = WhenMissing Identity + +newtype WhenMatched f k x y z = WhenMatched + { matchedKey :: forall a. k a -> x a -> y a -> f (Maybe (z a)) + } + +-- | Along with zipWithMaybeAMatched, witnesses the isomorphism between +-- @WhenMatched f k x y z@ and @forall a. k a -> x a -> y a -> f (Maybe (z a))@. +-- +-- @since UNRELEASED +runWhenMatched :: WhenMatched f k x y z -> k a -> x a -> y a -> f (Maybe (z a)) +runWhenMatched = matchedKey +{-# INLINE runWhenMatched #-} + +-- | Along with traverseMaybeMissing, witnesses the isomorphism between +-- @WhenMissing f k x y@ and @forall a. k a -> x a -> f (Maybe (y a))@. +-- +-- @since UNRELEASED +runWhenMissing :: WhenMissing f k x y -> k a -> x a -> f (Maybe (y a)) +runWhenMissing = missingKey +{-# INLINE runWhenMissing #-} + +-- | Map covariantly over a @'WhenMatched' f k x y@. +-- +-- @since UNRELEASED +mapWhenMatched + :: Functor f + => (forall a. g a -> h a) + -> WhenMatched f k x y g + -> WhenMatched f k x y h +mapWhenMatched f (WhenMatched g) = WhenMatched $ \k x y -> fmap (fmap f) (g k x y) +{-# INLINE mapWhenMatched #-} + +-- | A tactic for dealing with keys present in both maps in 'merge'. +-- +-- A tactic of type @ SimpleWhenMissing k x z @ is an abstract representation +-- of a function of type @ forall a. k a -> x a -> Maybe (z a) @. +-- +-- @since UNRELEASED +type SimpleWhenMatched = WhenMatched Identity + +-- | When a key is found in both maps, apply a function to the +-- key and values and use the result in the merged map. +-- +-- @ +-- zipWithMatched :: (forall a. k a -> x a -> y a -> z a) +-- -> SimpleWhenMatched k x y z +-- @ +-- +-- @since UNRELEASED +zipWithMatched + :: Applicative f + => (forall a. k a -> x a -> y a -> z a) + -> WhenMatched f k x y z +zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y +{-# INLINE zipWithMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values to produce an action and use its result in the merged map. +-- +-- @since UNRELEASED +zipWithAMatched + :: Applicative f + => (forall a. k a -> x a -> y a -> f (z a)) + -> WhenMatched f k x y z +zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y +{-# INLINE zipWithAMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values and maybe use the result in the merged map. +-- +-- @ +-- zipWithMaybeMatched :: (forall a. k a -> x a -> y a -> Maybe (z a)) +-- -> SimpleWhenMatched k x y z +-- @ +-- +-- @since UNRELEASED +zipWithMaybeMatched + :: Applicative f + => (forall a. k a -> x a -> y a -> Maybe (z a)) + -> WhenMatched f k x y z +zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y +{-# INLINE zipWithMaybeMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values, perform the resulting action, and maybe use +-- the result in the merged map. +-- +-- This is the fundamental 'WhenMatched' tactic. +-- +-- @since UNRELEASED +zipWithMaybeAMatched + :: (forall a. k a -> x a -> y a -> f (Maybe (z a))) + -> WhenMatched f k x y z +zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y +{-# INLINE zipWithMaybeAMatched #-} + + +-- | Drop all the entries whose keys are missing from the other +-- map. +-- +-- @ +-- dropMissing :: SimpleWhenMissing k x y +-- @ +-- +-- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing) +-- +-- but @dropMissing@ is much faster. +-- +-- @since UNRELEASED +dropMissing :: Applicative f => WhenMissing f k x y +dropMissing = WhenMissing + { missingSubtree = const (pure Tip) + , missingKey = \_ _ -> pure Nothing } +{-# INLINE dropMissing #-} + +-- | Preserve, unchanged, the entries whose keys are missing from +-- the other map. +-- +-- @ +-- preserveMissing :: SimpleWhenMissing k x x +-- @ +-- +-- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x) +-- +-- but @preserveMissing@ is much faster. +-- +-- @since UNRELEASED +preserveMissing :: Applicative f => WhenMissing f k x x +preserveMissing = WhenMissing + { missingSubtree = pure + , missingKey = \_ v -> pure (Just v) } +{-# INLINE preserveMissing #-} + +-- | Map over the entries whose keys are missing from the other map. +-- +-- @ +-- mapMissing :: (forall a. k a -> x a -> y a) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +-- +-- @since UNRELEASED +mapMissing :: Applicative f => (forall a. k a -> x a -> y a) -> WhenMissing f k x y +mapMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapWithKey f m + , missingKey = \ k x -> pure $ Just (f k x) } +{-# INLINE mapMissing #-} + +-- | Map over the entries whose keys are missing from the other map, +-- optionally removing some. This is the most powerful 'SimpleWhenMissing' +-- tactic, but others are usually more efficient. +-- +-- @ +-- mapMaybeMissing :: (forall a. k a -> x a -> Maybe (y a)) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) +-- +-- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations. +-- +-- @since UNRELEASED +mapMaybeMissing :: (GCompare k, Applicative f) => (forall a. k a -> x a -> Maybe (y a)) -> WhenMissing f k x y +mapMaybeMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapMaybeWithKey f m + , missingKey = \k x -> pure $! f k x } +{-# INLINE mapMaybeMissing #-} + +-- | Filter the entries whose keys are missing from the other map. +-- +-- @ +-- filterMissing :: (forall a. k a -> x a -> Bool) -> SimpleWhenMissing k x x +-- @ +-- +-- prop> filterMissing f = Merge.Lazy.mapMaybeMissing $ \k x -> guard (f k x) *> Just x +-- +-- but this should be a little faster. +-- +-- @since UNRELEASED +filterMissing + :: (GCompare k, Applicative f) + => (forall a. k a -> x a -> Bool) -> WhenMissing f k x x +filterMissing f = WhenMissing + { missingSubtree = \m -> pure $! filterWithKey f m + , missingKey = \k x -> pure $! if f k x then Just x else Nothing } +{-# INLINE filterMissing #-} + +-- | Filter the entries whose keys are missing from the other map +-- using some 'Applicative' action. +-- +-- @ +-- filterAMissing f = Merge.Lazy.traverseMaybeMissing $ +-- \k x -> (\b -> guard b *> Just x) <$> f k x +-- @ +-- +-- but this should be a little faster. +-- +-- @since UNRELEASED +filterAMissing + :: (GCompare k, Applicative f) + => (forall a. k a -> x a -> f Bool) -> WhenMissing f k x x +filterAMissing f = WhenMissing + { missingSubtree = \m -> filterWithKeyA f m + , missingKey = \k x -> bool Nothing (Just x) <$> f k x } +{-# INLINE filterAMissing #-} + +-- | This wasn't in Data.Bool until 4.7.0, so we define it here +bool :: a -> a -> Bool -> a +bool f _ False = f +bool _ t True = t + +-- | Traverse over the entries whose keys are missing from the other map. +-- +-- @since UNRELEASED +traverseMissing + :: Applicative f + => (forall a. k a -> x a -> f (y a)) -> WhenMissing f k x y +traverseMissing f = WhenMissing + { missingSubtree = traverseWithKey f + , missingKey = \k x -> Just <$> f k x } +{-# INLINE traverseMissing #-} + +-- | Traverse over the entries whose keys are missing from the other map, +-- optionally producing values to put in the result. +-- This is the most powerful 'WhenMissing' tactic, but others are usually +-- more efficient. +-- +-- @since UNRELEASED +traverseMaybeMissing + :: (GCompare k, Applicative f) + => (forall a. k a -> x a -> f (Maybe (y a))) -> WhenMissing f k x y +traverseMaybeMissing f = WhenMissing + { missingSubtree = traverseMaybeWithKey f + , missingKey = f } +{-# INLINE traverseMaybeMissing #-} + +merge :: GCompare k + => SimpleWhenMissing k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> SimpleWhenMissing k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> SimpleWhenMatched k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> DMap k a -- ^ DMap @m1@ + -> DMap k b -- ^ DMap @m2@ + -> DMap k c +merge g1 g2 f m1 m2 = runIdentity $ mergeA g1 g2 f m1 m2 + +mergeA + :: (Applicative f, GCompare k) + => WhenMissing f k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> WhenMissing f k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> WhenMatched f k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> DMap k a -- ^ DMap @m1@ + -> DMap k b -- ^ DMap @m2@ + -> f (DMap k c) +mergeA + WhenMissing{missingSubtree = g1t, missingKey = g1k} + WhenMissing{missingSubtree = g2t} + (WhenMatched f) = go + where + go t1 Tip = g1t t1 + go Tip t2 = g2t t2 + go (Bin _ kx x1 l1 r1) t2 = case splitLookup kx t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> liftA3 (\l' mx' r' -> maybe link2 (combine kx) mx' l' r') + l1l2 (g1k kx x1) r1r2 + Just x2 -> liftA3 (\l' mx' r' -> maybe link2 (combine kx) mx' l' r') + l1l2 (f kx x1 x2) r1r2 + where + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 +{-# INLINE mergeA #-}