{-# OPTIONS_HADDOCK not-home #-}

-- |
-- = WARNING
--
-- This module is considered __internal__.
--
-- The Package Versioning Policy __does not apply__.
--
-- The contents of this module may change __in any way whatsoever__
-- and __without any warning__ between minor versions of this package.
--
-- Authors importing this module are expected to track development
-- closely.
--
-- = Description
--
-- This module defines common constructs used by both "Data.Map.Merge.Set.Lazy"
-- and "Data.Map.Merge.Set.Strict".
--
-- @since 0.8.1
--
module Data.Map.Merge.Set.Internal
  ( WhenMatched(..)
  , SimpleWhenMatched
  , dropMatched
  , filterMatched
  , filterAMatched

  , WhenMissingSet(..)
  , SimpleWhenMissingSet
  , dropMissingSet

  , merge
  , mergeA

  , runWhenMatched
  , runWhenMissingSet
  ) where

import Control.Applicative (liftA3)
import Data.Functor.Identity (Identity(..))
import Data.Set (Set)
import qualified Data.Set.Internal as S
import Data.Map (Map)
import qualified Data.Map.Internal as M

-- | A tactic for dealing with keys present in both the set and the map in
-- 'merge' or 'mergeA'.
--
-- A tactic of type @WhenMatched f k a b@ is an abstract representation of
-- a function of type @k -> a -> f (Maybe b)@.
--
-- @since 0.8.1
newtype WhenMatched f k a b = WhenMatched
  { forall (f :: * -> *) k a b.
WhenMatched f k a b -> k -> a -> f (Maybe b)
matchedKey :: k -> a -> f (Maybe b)
  }

-- | Run @WhenMatched@.
--
-- @since 0.8.1
runWhenMatched :: WhenMatched f k a b -> k -> a -> f (Maybe b)
runWhenMatched :: forall (f :: * -> *) k a b.
WhenMatched f k a b -> k -> a -> f (Maybe b)
runWhenMatched = WhenMatched f k a b -> k -> a -> f (Maybe b)
forall (f :: * -> *) k a b.
WhenMatched f k a b -> k -> a -> f (Maybe b)
matchedKey

-- | A tactic for dealing with keys present in both the set and the map in
-- 'merge'.
--
-- A tactic of type @SimpleWhenMatched k a b@ is an abstract representation of
-- a function of type @k -> a -> Maybe b@.
--
-- @since 0.8.1
type SimpleWhenMatched = WhenMatched Identity

-- | When a key is found in both the map and the set, drop the key and value.
--
-- @since 0.8.1
dropMatched :: Applicative f => WhenMatched f k a b
dropMatched :: forall (f :: * -> *) k a b. Applicative f => WhenMatched f k a b
dropMatched = (k -> a -> f (Maybe b)) -> WhenMatched f k a b
forall (f :: * -> *) k a b.
(k -> a -> f (Maybe b)) -> WhenMatched f k a b
WhenMatched (\k
_ a
_ -> Maybe b -> f (Maybe b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe b
forall a. Maybe a
Nothing)
{-# INLINE dropMatched #-}

-- | When a key is found in both the map and the set, apply a function to the
-- key and the value in the map and keep the value in the merged map if the
-- result is @True@.
--
-- @since 0.8.1
filterMatched :: Applicative f => (k -> a -> Bool) -> WhenMatched f k a a
filterMatched :: forall (f :: * -> *) k a.
Applicative f =>
(k -> a -> Bool) -> WhenMatched f k a a
filterMatched k -> a -> Bool
f =
  (k -> a -> f (Maybe a)) -> WhenMatched f k a a
forall (f :: * -> *) k a b.
(k -> a -> f (Maybe b)) -> WhenMatched f k a b
WhenMatched (\k
k a
x -> if k -> a -> Bool
f k
k a
x then Maybe a -> f (Maybe a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
x) else Maybe a -> f (Maybe a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing)
{-# INLINE filterMatched #-}

-- | When a key is found in both the map and the set, apply a function to the
-- key and the value in the map and keep the value in the merged map if the
-- result of the action is @True@.
--
-- @since 0.8.1
filterAMatched :: Functor f => (k -> a -> f Bool) -> WhenMatched f k a a
filterAMatched :: forall (f :: * -> *) k a.
Functor f =>
(k -> a -> f Bool) -> WhenMatched f k a a
filterAMatched k -> a -> f Bool
f =
  (k -> a -> f (Maybe a)) -> WhenMatched f k a a
forall (f :: * -> *) k a b.
(k -> a -> f (Maybe b)) -> WhenMatched f k a b
WhenMatched (\k
k a
x -> (\Bool
b -> if Bool
b then a -> Maybe a
forall a. a -> Maybe a
Just a
x else Maybe a
forall a. Maybe a
Nothing) (Bool -> Maybe a) -> f Bool -> f (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> a -> f Bool
f k
k a
x)
{-# INLINE filterAMatched #-}

-- | A tactic for dealing with keys present in the set but not in the map in
-- 'merge' or 'mergeA'.
--
-- A tactic of type @WhenMissingSet f k a@ is an abstract representation of
-- a function of type @k -> f (Maybe a)@.
--
-- @since 0.8.1
data WhenMissingSet f k a = WhenMissingSet
  { forall (f :: * -> *) k a.
WhenMissingSet f k a -> Set k -> f (Map k a)
missingSubtree :: Set k -> f (Map k a)
  , forall (f :: * -> *) k a. WhenMissingSet f k a -> k -> f (Maybe a)
missingKey :: k -> f (Maybe a)
  }

-- | Run @WhenMissingSet@.
--
-- @since 0.8.1
runWhenMissingSet :: WhenMissingSet f k a -> k -> f (Maybe a)
runWhenMissingSet :: forall (f :: * -> *) k a. WhenMissingSet f k a -> k -> f (Maybe a)
runWhenMissingSet = WhenMissingSet f k a -> k -> f (Maybe a)
forall (f :: * -> *) k a. WhenMissingSet f k a -> k -> f (Maybe a)
missingKey

-- | A tactic for dealing with keys present in the set but not in the map in
-- 'merge'.
--
-- A tactic of type @SimpleWhenMissingSet k a@ is an abstract representation of
-- a function of type @k -> Maybe a@.
--
-- @since 0.8.1
type SimpleWhenMissingSet = WhenMissingSet Identity

-- | Drop keys that are present in the set but missing from the map.
--
-- @since 0.8.1
dropMissingSet :: Applicative f => WhenMissingSet f k a
dropMissingSet :: forall (f :: * -> *) k a. Applicative f => WhenMissingSet f k a
dropMissingSet = WhenMissingSet
  { missingSubtree :: Set k -> f (Map k a)
missingSubtree = \Set k
_ -> Map k a -> f (Map k a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k a
forall k a. Map k a
M.empty
  , missingKey :: k -> f (Maybe a)
missingKey = \k
_ -> Maybe a -> f (Maybe a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
  }
{-# INLINE dropMissingSet #-}

-- | Merge a map and a set into a map.
--
-- 'merge' takes a 'M.SimpleWhenMissing' tactic, a 'SimpleWhenMissingSet'
-- tactic, a 'SimpleWhenMatched' tactic, a map and a set. It uses the tactics to
-- merge the map and the set into a map.
--
-- Its behavior is best understood via the tactics @mapMaybeMissing@,
-- @generateMaybeMissingSet@, and @mapMaybeMatched@. Consider
--
-- @
-- merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2
-- @
--
-- @
-- g1 k x = if k == 2 then Just ("1" ++ x) else Nothing
-- g2 k = if k == 3 then Just "2" else Nothing
-- f k x = if k == 6 then Just ("3" ++ x) else Nothing
-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")]
-- s2 = fromList [3, 6, 9, 12]
-- @
--
-- 'merge' will pass the keys and values to @g1@, @g2@, or @f@ as appropriate,
-- producing a @Maybe@ for each element.
--
-- @
-- m1:      [ (2, "a"),           (4, "b"),  (6, "c"), (8, "d"),          (10, "e"), (12, "f")]
-- s2:      [                  3,                   6,                 9,                   12]
-- result:  [ g1 2 "a",     g2 3, g1 4 "b",   f 6 "c", g1 8 "d",    g2 9, g1 10 "e",  f 12 "f"]
--        = [Just "1a", Just "2",  Nothing, Just "3c",  Nothing, Nothing,   Nothing,   Nothing]
-- @
--
-- The result map contains the @Just@ values.
--
-- >>> merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2
-- fromList [(2,"1a"), (3,"2g"), (6,"3c")]
--
-- When 'merge' is given three arguments, it is inlined at the call
-- site. To prevent excessive inlining, you should typically use 'merge'
-- to define your custom combining functions.
--
-- @since 0.8.1
merge
  :: Ord k
  => M.SimpleWhenMissing k a b -- ^ What to do with keys in @m1@ but not @s2@
  -> SimpleWhenMissingSet k b -- ^ What to do with keys in @s2@ but not @m1@
  -> SimpleWhenMatched k a b -- ^ What to do with keys in both @m1@ and @s2@
  -> Map k a -- ^ Map @m1@
  -> Set k -- ^ Set @s2@
  -> Map k b
merge :: forall k a b.
Ord k =>
SimpleWhenMissing k a b
-> SimpleWhenMissingSet k b
-> SimpleWhenMatched k a b
-> Map k a
-> Set k
-> Map k b
merge SimpleWhenMissing k a b
miss1 SimpleWhenMissingSet k b
miss2 SimpleWhenMatched k a b
match = \Map k a
t1 Set k
t2 -> Identity (Map k b) -> Map k b
forall a. Identity a -> a
runIdentity (SimpleWhenMissing k a b
-> SimpleWhenMissingSet k b
-> SimpleWhenMatched k a b
-> Map k a
-> Set k
-> Identity (Map k b)
forall (f :: * -> *) k a b.
(Applicative f, Ord k) =>
WhenMissing f k a b
-> WhenMissingSet f k b
-> WhenMatched f k a b
-> Map k a
-> Set k
-> f (Map k b)
mergeA SimpleWhenMissing k a b
miss1 SimpleWhenMissingSet k b
miss2 SimpleWhenMatched k a b
match Map k a
t1 Set k
t2)
{-# INLINE merge #-}

-- | Merge a map and a set into a map. Applicative version of 'merge'.
--
-- 'mergeA' takes a 'M.WhenMissing' tactic, a 'WhenMissingSet' tactic, a
-- 'WhenMatched' tactic, a map and a set. It uses the tactics to merge the map
-- and the set into a map.
--
-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are
-- performed in increasing order of keys.
--
-- Consider
--
-- @
-- mergeA (traverseMaybeMissing g1)
--        (generateMaybeAMissingSet g2)
--        (traverseMaybeMatched f)
--        m1
--        s2
-- @
--
-- @
-- g1 k x = let z = if k == 2 then Just ("1" ++ x) else Nothing
--          in z <$ putStrLn ("g1 " ++ show (k, x))
-- g2 k = let z = if k == 3 then Just "2" else Nothing
--        in z <$ putStrLn ("g2 " ++ show k)
-- f k x = let z = if k == 6 then Just ("3" ++ x) else Nothing
--         in z <$ putStrLn ("f " ++ show (k, x))
-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")]
-- m2 = fromList [3, 6, 9, 12]
-- @
--
-- As with 'merge', the result map is @[(2,"1a"), (3,"2"), (6,"3c")]@.
-- Additionally, @g1@, @g2@, and @f@ perform @IO@ effects, printing in
-- increasing order of key.
--
-- >>> mergeA (traverseMaybeMissing g1) (generateMaybeAMissingSet g2) (traverseMaybeMatched f) m1 s2
-- g1 (2,"a")
-- g2 3
-- g1 (4,"b")
-- f (6,"c")
-- g1 (8,"d")
-- g2 9
-- g1 (10,"e")
-- f (12,"f")
-- fromList [(2,"1a"),(3,"2"),(6,"3c")]
--
-- When 'mergeA' is given three arguments, it is inlined at the call
-- site. To prevent excessive inlining, you should generally only use
-- 'mergeA' to define custom combining functions.
--
-- === __Examples__
--
-- @
-- data Pair a = Pair !a !a deriving Functor
--
-- instance Applicative Pair where
--    pure x = Pair x x
--    liftA2 f (Pair x1 y1) (Pair x2 y2) = Pair (f x1 x2) (f y1 y2)
--
-- -- | Partition the map according to whether the keys appear in the set.
-- partitionKeys :: Ord k => Map k a -> Set k -> (Map k a, Map k a)
-- partitionKeys m s =
--   case mergeA dropAndPreserveMissing dropMissingSet preserveAndDropMatched m s of
--     Pair m1 m2 -> (m1, m2)
--   where
--     dropAndPreserveMissing = whenMissing (\\_k x -> Pair Nothing (Just x)) (\\m -> Pair empty m)
--     preserveAndDropMatched = traverseMaybeMatched (\\_k x -> Pair (Just x) Nothing)
-- @
--
-- @
-- import Data.Functor.Const (Const(..))
-- import Data.Monoid (All(..))
--
-- -- | Whether the keys of the map are a subset of the keys of the set.
-- keysAreSubsetOf :: Ord k => Map k a -> Set k -> Bool
-- keysAreSubsetOf m s =
--   getAll (getConst (mergeA isEmpty dropMissing 'dropMatched' m1 m2))
--   where
--     isEmpty = whenMissing (\\_k _x -> Const (All False)) (\\m -> Const (All (null m)))
-- @
--
-- @since 0.8.1
mergeA
  :: (Applicative f, Ord k)
  => M.WhenMissing f k a b -- ^ What to do with keys in @m1@ but not @s2@
  -> WhenMissingSet f k b -- ^ What to do with keys in @s2@ but not @m1@
  -> WhenMatched f k a b -- ^ What to do with keys in both @m1@ and @s2@
  -> Map k a -- ^ Map @m1@
  -> Set k -- ^ Set @s2@
  -> f (Map k b)
mergeA :: forall (f :: * -> *) k a b.
(Applicative f, Ord k) =>
WhenMissing f k a b
-> WhenMissingSet f k b
-> WhenMatched f k a b
-> Map k a
-> Set k
-> f (Map k b)
mergeA
  M.WhenMissing{missingSubtree :: forall (f :: * -> *) k x y.
WhenMissing f k x y -> Map k x -> f (Map k y)
M.missingSubtree = Map k a -> f (Map k b)
g1t, missingKey :: forall (f :: * -> *) k x y.
WhenMissing f k x y -> k -> x -> f (Maybe y)
M.missingKey = k -> a -> f (Maybe b)
g1k}
  WhenMissingSet{missingSubtree :: forall (f :: * -> *) k a.
WhenMissingSet f k a -> Set k -> f (Map k a)
missingSubtree = Set k -> f (Map k b)
g2t}
  WhenMatched{matchedKey :: forall (f :: * -> *) k a b.
WhenMatched f k a b -> k -> a -> f (Maybe b)
matchedKey = k -> a -> f (Maybe b)
f} = Map k a -> Set k -> f (Map k b)
go
  where
    go :: Map k a -> Set k -> f (Map k b)
go Map k a
t1 Set k
S.Tip = Map k a -> f (Map k b)
g1t Map k a
t1
    go Map k a
M.Tip Set k
t2 = Set k -> f (Map k b)
g2t Set k
t2
    go (M.Bin Size
_ k
k1 a
x1 Map k a
l1 Map k a
r1) Set k
t2 = case k -> Set k -> (Set k, Bool, Set k)
forall a. Ord a => a -> Set a -> (Set a, Bool, Set a)
S.splitMember k
k1 Set k
t2 of
      (Set k
l2, Bool
found, Set k
r2) ->
        (Map k b -> Maybe b -> Map k b -> Map k b)
-> f (Map k b) -> f (Maybe b) -> f (Map k b) -> f (Map k b)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3
          (\Map k b
l' Maybe b
mx' Map k b
r' -> (Map k b -> Map k b -> Map k b)
-> (b -> Map k b -> Map k b -> Map k b)
-> Maybe b
-> Map k b
-> Map k b
-> Map k b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map k b -> Map k b -> Map k b
forall k a. Map k a -> Map k a -> Map k a
M.link2 (k -> b -> Map k b -> Map k b -> Map k b
forall k a. k -> a -> Map k a -> Map k a -> Map k a
M.link k
k1) Maybe b
mx' Map k b
l' Map k b
r')
          (Map k a -> Set k -> f (Map k b)
go Map k a
l1 Set k
l2)
          (if Bool
found then k -> a -> f (Maybe b)
f k
k1 a
x1 else k -> a -> f (Maybe b)
g1k k
k1 a
x1)
          (Map k a -> Set k -> f (Map k b)
go Map k a
r1 Set k
r2)
{-# INLINE mergeA #-}