Add filterKeys for Map and IntMap

containers-tests/tests/intmap-properties.hs

@@ -22,7 +22,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
@@ -102,7 +102,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
@@ -179,6 +180,9 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "deleteMin"            prop_deleteMinModel
              , testProperty "deleteMax"            prop_deleteMaxModel
              , testProperty "filter"               prop_filter
+             , testProperty "filterWithKey"        prop_filterWithKey
+             , testProperty "filterKeys"           prop_filterKeys
+             , testProperty "filterKeysFidelity"   prop_filterKeysFidelity
              , testProperty "partition"            prop_partition
              , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
              , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
@@ -858,8 +862,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"
 
@@ -1464,6 +1473,24 @@ prop_filter p ys = length ys > 0 ==>
   in  valid m .&&.
       m === fromList (List.filter (apply p . snd) xs)
 
+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_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_filterKeysFidelity :: Fun Int Bool -> IMap -> Property
+prop_filterKeysFidelity p m = fwk === fk
+  where
+    fwk = filterWithKey (\k _ -> apply p k) m
+    fk = filterKeys (apply p) 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

@@ -115,7 +115,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
@@ -820,8 +821,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

containers/src/Data/IntMap/Internal.hs

@@ -223,6 +223,7 @@ module Data.IntMap.Internal (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
     , restrictKeys
     , withoutKeys
@@ -2578,6 +2579,14 @@ 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 (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
+-- > filterKeys (> 4) == filterWithKey (\k _ -> k > 4)
+
+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

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

containers/src/Data/IntMap/Strict.hs

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

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

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

containers/src/Data/Map/Internal.hs

@@ -288,6 +288,7 @@ module Data.Map.Internal (
 
     -- * Filter
     , filter
+    , filterKeys
     , filterWithKey
 
     , takeWhileAntitone
@@ -2857,6 +2858,14 @@ 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 (> 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

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

containers/src/Data/Map/Strict.hs

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

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

@@ -250,6 +250,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