Add filterKeys for Map and IntMap

containers-tests/tests/intmap-properties.hs

@@ -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
@@ -858,8 +859,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"
 

containers-tests/tests/map-properties.hs

@@ -115,7 +115,7 @@ main = defaultMain $ testGroup "map-properties"
          , testCase "fromDistinctAscList" test_fromDistinctAscList
          , testCase "fromDistinctDescList" test_fromDistinctDescList
          , testCase "filter" test_filter
-         , testCase "filterWithKey" test_filteWithKey
+         , testCase "filterWithKey" test_filterWithKey
          , testCase "partition" test_partition
          , testCase "partitionWithKey" test_partitionWithKey
          , testCase "mapMaybe" test_mapMaybe
@@ -820,8 +820,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,17 @@ 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 :: (Key -> Bool) -> IntMap a -> IntMap a
+filterKeys predicate = go
+    where
+    go Nil           = Nil
+    go t@(Tip k _)   = if predicate k then t else Nil
+    go (Bin p m l r) = bin p m (go l) (go r)
+
 -- | \(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,13 @@ 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"
+
+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