Add filterKeys for Map and IntMap (#972)

containers-tests/tests/intmap-properties.hs

@@ -23,7 +23,7 @@ import Data.Foldable (foldMap)
 import Data.Function
 import Data.Traversable (Traversable(traverse), foldMapDefault)
 import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl')
-import qualified Prelude (map)
+import qualified Prelude (map, filter)
 
 import Data.List (nub,sort)
 import qualified Data.List as List
@@ -103,7 +103,8 @@ main = defaultMain $ testGroup "intmap-properties"
              , testCase "fromAscListWithKey" test_fromAscListWithKey
              , testCase "fromDistinctAscList" test_fromDistinctAscList
              , testCase "filter" test_filter
-             , testCase "filterWithKey" test_filteWithKey
+             , testCase "filterKeys" test_filterKeys
+             , testCase "filterWithKey" test_filterWithKey
              , testCase "partition" test_partition
              , testCase "partitionWithKey" test_partitionWithKey
              , testCase "mapMaybe" test_mapMaybe
@@ -182,6 +183,8 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "deleteMin"            prop_deleteMinModel
              , testProperty "deleteMax"            prop_deleteMaxModel
              , testProperty "filter"               prop_filter
+             , testProperty "filterKeys"           prop_filterKeys
+             , testProperty "filterWithKey"        prop_filterWithKey
              , testProperty "partition"            prop_partition
              , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
              , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
@@ -867,8 +870,13 @@ test_filter = do
     filter (> "x") (fromList [(5,"a"), (-3,"b")]) @?= empty
     filter (< "a") (fromList [(5,"a"), (-3,"b")]) @?= empty
 
-test_filteWithKey :: Assertion
-test_filteWithKey = do
+test_filterKeys :: Assertion
+test_filterKeys = do
+    filterKeys (> 4) (fromList [(5,"a"), (3,"b")])  @?= singleton 5 "a"
+    filterKeys (> 4) (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "a"
+
+test_filterWithKey :: Assertion
+test_filterWithKey = do
     filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")])  @?= singleton 5 "a"
     filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "a"
 
@@ -1506,6 +1514,18 @@ prop_filter p ys = length ys > 0 ==>
   in  valid m .&&.
       m === fromList (List.filter (apply p . snd) xs)
 
+prop_filterKeys :: Fun Int Bool -> IMap -> Property
+prop_filterKeys fun m =
+  valid m' .&&. toList m' === Prelude.filter (apply fun . fst) (toList m)
+  where
+    m' = filterKeys (apply fun) m
+
+prop_filterWithKey :: Fun (Int, Int) Bool -> IMap -> Property
+prop_filterWithKey fun m =
+  valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
+  where
+    m' = filterWithKey (applyFun2 fun) m
+
 prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
 prop_partition p ys = length ys > 0 ==>
   let xs = List.nubBy ((==) `on` fst) ys

containers-tests/tests/map-properties.hs

@@ -118,7 +118,8 @@ main = defaultMain $ testGroup "map-properties"
          , testCase "fromDistinctAscList" test_fromDistinctAscList
          , testCase "fromDistinctDescList" test_fromDistinctDescList
          , testCase "filter" test_filter
-         , testCase "filterWithKey" test_filteWithKey
+         , testCase "filterKeys" test_filterKeys
+         , testCase "filterWithKey" test_filterWithKey
          , testCase "partition" test_partition
          , testCase "partitionWithKey" test_partitionWithKey
          , testCase "mapMaybe" test_mapMaybe
@@ -223,6 +224,8 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "deleteMin"            prop_deleteMinModel
          , testProperty "deleteMax"            prop_deleteMaxModel
          , testProperty "filter"               prop_filter
+         , testProperty "filterKeys"           prop_filterKeys
+         , testProperty "filterWithKey"        prop_filterWithKey
          , testProperty "partition"            prop_partition
          , testProperty "map"                  prop_map
          , testProperty "fmap"                 prop_fmap
@@ -796,8 +799,11 @@ test_filter = do
     filter (> "x") (fromList [(5,"a"), (3,"b")]) @?= empty
     filter (< "a") (fromList [(5,"a"), (3,"b")]) @?= empty
 
-test_filteWithKey :: Assertion
-test_filteWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
+test_filterKeys :: Assertion
+test_filterKeys = filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
+
+test_filterWithKey :: Assertion
+test_filterWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
 
 test_partition :: Assertion
 test_partition = do
@@ -1463,6 +1469,18 @@ prop_filter p ys = length ys > 0 ==>
       m  = fromList xs
   in  filter (apply p) m == fromList (List.filter (apply p . snd) xs)
 
+prop_filterKeys :: Fun Int Bool -> IMap -> Property
+prop_filterKeys fun m =
+  valid m' .&&. toList m' === Prelude.filter (apply fun . fst) (toList m)
+  where
+    m' = filterKeys (apply fun) m
+
+prop_filterWithKey :: Fun (Int, Int) Bool -> IMap -> Property
+prop_filterWithKey fun m =
+  valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
+  where
+    m' = filterWithKey (apply2 fun) m
+
 prop_take :: Int -> Map Int Int -> Property
 prop_take n xs = valid taken .&&.
                  taken === fromDistinctAscList (List.take n (toList xs))

containers/src/Data/IntMap/Internal.hs

@@ -218,6 +218,7 @@ module Data.IntMap.Internal (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys
@@ -2638,6 +2639,19 @@ filter :: (a -> Bool) -> IntMap a -> IntMap a
 filter p m
   = filterWithKey (\_ x -> p x) m
 
+-- | \(O(n)\). Filter all keys that satisfy some predicate.
+--
+-- @
+-- filterKeys p = 'filterWithKey' (\\k _ -> p k)
+-- @
+--
+-- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
+--
+-- @since FIXME
+
+filterKeys :: (Key -> Bool) -> IntMap a -> IntMap a
+filterKeys predicate = filterWithKey (\k _ -> predicate k)
+
 -- | \(O(n)\). Filter all keys\/values that satisfy some predicate.
 --
 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"

containers/src/Data/IntMap/Lazy.hs

@@ -198,6 +198,7 @@ module Data.IntMap.Lazy (
 
     -- * Filter
     , IM.filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys

containers/src/Data/IntMap/Strict.hs

@@ -216,6 +216,7 @@ module Data.IntMap.Strict (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys

containers/src/Data/IntMap/Strict/Internal.hs

@@ -214,6 +214,7 @@ module Data.IntMap.Strict.Internal (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys
@@ -279,6 +280,7 @@ import Data.IntMap.Internal
   , empty
   , assocs
   , filter
+  , filterKeys
   , filterWithKey
   , findMin
   , findMax

containers/src/Data/Map/Internal.hs

@@ -291,6 +291,7 @@ module Data.Map.Internal (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
 
     , takeWhileAntitone
@@ -2944,6 +2945,19 @@ filter :: (a -> Bool) -> Map k a -> Map k a
 filter p m
   = filterWithKey (\_ x -> p x) m
 
+-- | \(O(n)\). Filter all keys that satisfy the predicate.
+--
+-- @
+-- filterKeys p = 'filterWithKey' (\\k _ -> p k)
+-- @
+--
+-- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
+--
+-- @since FIXME
+
+filterKeys :: (k -> Bool) -> Map k a -> Map k a
+filterKeys p m = filterWithKey (\k _ -> p k) m
+
 -- | \(O(n)\). Filter all keys\/values that satisfy the predicate.
 --
 -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"

containers/src/Data/Map/Lazy.hs

@@ -228,6 +228,7 @@ module Data.Map.Lazy (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys

containers/src/Data/Map/Strict.hs

@@ -243,6 +243,7 @@ module Data.Map.Strict
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys

containers/src/Data/Map/Strict/Internal.hs

@@ -253,6 +253,7 @@ module Data.Map.Strict.Internal
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys
@@ -359,6 +360,7 @@ import Data.Map.Internal
   , drop
   , dropWhileAntitone
   , filter
+  , filterKeys
   , filterWithKey
   , findIndex
   , findMax