Add 'compose' for maps (#672)

Add 'compose' for maps

Author: alexfmpe

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

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

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
 --------------------------------------------------------------------}

containers/src/Data/IntMap/Lazy.hs

@@ -146,6 +146,9 @@ module Data.IntMap.Lazy (
     -- ** Disjoint
     , disjoint
 
+    -- ** Compose
+    , compose
+
     -- ** Universal combining function
     , mergeWithKey
 

containers/src/Data/IntMap/Strict.hs

@@ -165,6 +165,9 @@ module Data.IntMap.Strict (
     -- ** Disjoint
     , disjoint
 
+    -- ** Compose
+    , compose
+
     -- ** Universal combining function
     , mergeWithKey
 

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
 --------------------------------------------------------------------}

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 }

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"
 

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"
 

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
 --------------------------------------------------------------------}