Internal.hs

{-# 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
  { matchedKey :: k -> a -> f (Maybe b)
  }

-- | Run @WhenMatched@.
--
-- @since 0.8.1
runWhenMatched :: WhenMatched f k a b -> k -> a -> f (Maybe b)
runWhenMatched = 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 = WhenMatched (\_ _ -> pure 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 f =
  WhenMatched (\k x -> if f k x then pure (Just x) else pure 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 f =
  WhenMatched (\k x -> (\b -> if b then Just x else Nothing) <$> f k 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
  { missingSubtree :: Set k -> f (Map k a)
  , missingKey :: k -> f (Maybe a)
  }

-- | Run @WhenMissingSet@.
--
-- @since 0.8.1
runWhenMissingSet :: WhenMissingSet f k a -> k -> f (Maybe a)
runWhenMissingSet = 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 = WhenMissingSet
  { missingSubtree = \_ -> pure M.empty
  , missingKey = \_ -> pure 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 miss1 miss2 match = \t1 t2 -> runIdentity (mergeA miss1 miss2 match t1 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
  M.WhenMissing{M.missingSubtree = g1t, M.missingKey = g1k}
  WhenMissingSet{missingSubtree = g2t}
  WhenMatched{matchedKey = f} = go
  where
    go t1 S.Tip = g1t t1
    go M.Tip t2 = g2t t2
    go (M.Bin _ k1 x1 l1 r1) t2 = case S.splitMember k1 t2 of
      (l2, found, r2) ->
        liftA3
          (\l' mx' r' -> maybe M.link2 (M.link k1) mx' l' r')
          (go l1 l2)
          (if found then f k1 x1 else g1k k1 x1)
          (go r1 r2)
{-# INLINE mergeA #-}