diff --git a/containers-tests/tests/intmap-properties.hs b/containers-tests/tests/intmap-properties.hs index 4a55bc0b9..aba4c2f12 100644 --- a/containers-tests/tests/intmap-properties.hs +++ b/containers-tests/tests/intmap-properties.hs @@ -13,6 +13,7 @@ import Data.IntMap.Internal.Debug (showTree) import IntMapValidity (valid) import Control.Applicative (Applicative(..)) +import Control.Monad ((<=<)) import Data.Monoid import Data.Maybe hiding (mapMaybe) import qualified Data.Maybe as Maybe (mapMaybe) @@ -174,6 +175,7 @@ main = defaultMain , testProperty "lookupLE" prop_lookupLE , testProperty "lookupGE" prop_lookupGE , testProperty "disjoint" prop_disjoint + , testProperty "compose" prop_compose , testProperty "lookupMin" prop_lookupMin , testProperty "lookupMax" prop_lookupMax , testProperty "findMin" prop_findMin @@ -1243,6 +1245,9 @@ prop_intersectionWithKeyModel xs ys prop_disjoint :: UMap -> UMap -> Property prop_disjoint m1 m2 = disjoint m1 m2 === null (intersection m1 m2) +prop_compose :: IMap -> IMap -> Int -> Property +prop_compose bc ab k = (compose bc ab !? k) === ((bc !?) <=< (ab !?)) k + -- TODO: the second argument should be simply an 'IntSet', but that -- runs afoul of our orphan instance. prop_restrictKeys :: IMap -> IMap -> Property diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs index 5665e10a6..cf7fe02fe 100644 --- a/containers-tests/tests/map-properties.hs +++ b/containers-tests/tests/map-properties.hs @@ -13,7 +13,7 @@ import Data.Map.Internal.Debug (showTree, showTreeWith, balanced) import Control.Applicative (Const(Const, getConst), pure, (<$>), (<*>)) import Control.Monad.Trans.State.Strict import Control.Monad.Trans.Class -import Control.Monad (liftM4) +import Control.Monad (liftM4, (<=<)) import Data.Functor.Identity (Identity(Identity, runIdentity)) import Data.Monoid import Data.Maybe hiding (mapMaybe) @@ -180,6 +180,7 @@ main = defaultMain , testProperty "intersectionWithKey" prop_intersectionWithKey , testProperty "intersectionWithKeyModel" prop_intersectionWithKeyModel , testProperty "disjoint" prop_disjoint + , testProperty "compose" prop_compose , testProperty "differenceMerge" prop_differenceMerge , testProperty "unionWithKeyMerge" prop_unionWithKeyMerge , testProperty "mergeWithKey model" prop_mergeWithKeyModel @@ -1160,6 +1161,9 @@ prop_intersectionWithKeyModel xs ys prop_disjoint :: UMap -> UMap -> Property prop_disjoint m1 m2 = disjoint m1 m2 === null (intersection m1 m2) +prop_compose :: IMap -> IMap -> Int -> Property +prop_compose bc ab k = (compose bc ab !? k) === ((bc !?) <=< (ab !?)) k + prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool prop_mergeWithKeyModel xs ys = and [ testMergeWithKey f keep_x keep_y diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs index b8bc3ffb8..38f7c8419 100644 --- a/containers/src/Data/IntMap/Internal.hs +++ b/containers/src/Data/IntMap/Internal.hs @@ -131,6 +131,9 @@ module Data.IntMap.Internal ( , intersectionWith , intersectionWithKey + -- ** Compose + , compose + -- ** General combining function , SimpleWhenMissing , SimpleWhenMatched @@ -765,6 +768,25 @@ disjoint t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) | zero p1 m2 = disjoint t1 l2 | otherwise = disjoint t1 r2 +{-------------------------------------------------------------------- + Compose +--------------------------------------------------------------------} +-- | /O(|ab|*min(|bc|,W))/. Relate the keys of one map to the values of +-- the other, by using the values of the former as keys for lookups +-- in the latter. +-- +-- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] +-- +-- @ +-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') +-- @ +-- +-- @since UNRELEASED +compose :: IntMap c -> IntMap Int -> IntMap c +compose bc !ab + | null bc = empty + | otherwise = mapMaybe (bc !?) ab + {-------------------------------------------------------------------- Construction --------------------------------------------------------------------} diff --git a/containers/src/Data/IntMap/Lazy.hs b/containers/src/Data/IntMap/Lazy.hs index 8d4ce8c82..3f38e3c29 100644 --- a/containers/src/Data/IntMap/Lazy.hs +++ b/containers/src/Data/IntMap/Lazy.hs @@ -146,6 +146,9 @@ module Data.IntMap.Lazy ( -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** Universal combining function , mergeWithKey diff --git a/containers/src/Data/IntMap/Strict.hs b/containers/src/Data/IntMap/Strict.hs index ef732c2e2..36bbad1e2 100644 --- a/containers/src/Data/IntMap/Strict.hs +++ b/containers/src/Data/IntMap/Strict.hs @@ -165,6 +165,9 @@ module Data.IntMap.Strict ( -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** Universal combining function , mergeWithKey diff --git a/containers/src/Data/IntMap/Strict/Internal.hs b/containers/src/Data/IntMap/Strict/Internal.hs index 2b695c0fe..639d1acba 100644 --- a/containers/src/Data/IntMap/Strict/Internal.hs +++ b/containers/src/Data/IntMap/Strict/Internal.hs @@ -163,6 +163,9 @@ module Data.IntMap.Strict.Internal ( -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** Universal combining function , mergeWithKey @@ -715,6 +718,25 @@ intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c intersectionWithKey f m1 m2 = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) (const Nil) (const Nil) m1 m2 +{-------------------------------------------------------------------- + Compose +--------------------------------------------------------------------} +-- | /O(|ab|*min(|bc|,W))/. Relate the keys of one map to the values of +-- the other, by using the values of the former as keys for lookups +-- in the latter. +-- +-- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] +-- +-- @ +-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') +-- @ +-- +-- @since UNRELEASED +compose :: IntMap c -> IntMap Int -> IntMap c +compose bc !ab + | null bc = empty + | otherwise = mapMaybe (bc !?) ab + {-------------------------------------------------------------------- MergeWithKey --------------------------------------------------------------------} diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index b9fa2eeca..40cc2d308 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -192,6 +192,9 @@ module Data.Map.Internal ( -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** General combining function , SimpleWhenMissing , SimpleWhenMatched @@ -2088,6 +2091,25 @@ disjoint (Bin _ k _ l r) t where (lt,found,gt) = splitMember k t +{-------------------------------------------------------------------- + Compose +--------------------------------------------------------------------} +-- | /O(|ab|*log(|bc|))/. Relate the keys of one map to the values of +-- the other, by using the values of the former as keys for lookups +-- in the latter. +-- +-- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] +-- +-- @ +-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') +-- @ +-- +-- @since UNRELEASED +compose :: Ord b => Map b c -> Map a b -> Map a c +compose bc !ab + | null bc = empty + | otherwise = mapMaybe (bc !?) ab + #if !MIN_VERSION_base (4,8,0) -- | The identity type. newtype Identity a = Identity { runIdentity :: a } diff --git a/containers/src/Data/Map/Lazy.hs b/containers/src/Data/Map/Lazy.hs index 37ae931e6..6b05603a1 100644 --- a/containers/src/Data/Map/Lazy.hs +++ b/containers/src/Data/Map/Lazy.hs @@ -166,6 +166,9 @@ module Data.Map.Lazy ( -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** General combining functions -- | See "Data.Map.Merge.Lazy" diff --git a/containers/src/Data/Map/Strict.hs b/containers/src/Data/Map/Strict.hs index a73078413..8eea8c329 100644 --- a/containers/src/Data/Map/Strict.hs +++ b/containers/src/Data/Map/Strict.hs @@ -182,6 +182,9 @@ module Data.Map.Strict -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** General combining functions -- | See "Data.Map.Merge.Strict" diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs index 06c76b8f6..ab1d11744 100644 --- a/containers/src/Data/Map/Strict/Internal.hs +++ b/containers/src/Data/Map/Strict/Internal.hs @@ -145,6 +145,9 @@ module Data.Map.Strict.Internal -- ** Disjoint , disjoint + -- ** Compose + , compose + -- ** General combining function , SimpleWhenMissing , SimpleWhenMatched @@ -1200,6 +1203,25 @@ forceMaybe Nothing = Nothing forceMaybe m@(Just !_) = m {-# INLINE forceMaybe #-} +{-------------------------------------------------------------------- + Compose +--------------------------------------------------------------------} +-- | /O(|ab|*log(|bc|))/. Relate the keys of one map to the values of +-- the other, by using the values of the former as keys for lookups +-- in the latter. +-- +-- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] +-- +-- @ +-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') +-- @ +-- +-- @since UNRELEASED +compose :: Ord b => Map b c -> Map a b -> Map a c +compose bc !ab + | null bc = empty + | otherwise = mapMaybe (bc !?) ab + {-------------------------------------------------------------------- MergeWithKey --------------------------------------------------------------------}