Changes after review

containers-tests/tests/intmap-strictness.hs

@@ -87,6 +87,24 @@ pInsertLookupWithKeyValueStrict f k v m
                      not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m)
     | otherwise    = isBottom $ M.insertLookupWithKey (apply3 f) k bottom m
 
+pFilterWithKey :: Fun (Int, Int) Bool -> IMap -> Property
+pFilterWithKey fun m =
+  valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
+  where
+    m' = filterWithKey (apply2 fun) m
+
+-- pFilterKeys :: Fun (Int, Int) Bool -> IMap -> Property
+-- pFilterKeys fun m =
+--   valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
+--   where
+--     m' = filterKeys (apply2 fun) m
+
+-- pFilter :: Fun (Int, Int) Bool -> IMap -> Property
+-- pFilter fun m =
+--   valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
+--   where
+--     m' = filter (apply2 fun) m
+
 ------------------------------------------------------------------------
 -- test a corner case of fromAscList
 --
@@ -199,6 +217,7 @@ tests =
         pInsertLookupWithKeyValueStrict
       , testProperty "fromAscList is somewhat value-lazy" pFromAscListLazy
       , testProperty "fromAscList is somewhat value-strict" pFromAscListStrict
+      , testProperty "filterWithKey" pFilterWithKey
 #if __GLASGOW_HASKELL__ >= 806
       , testProperty "strict foldr'" pStrictFoldr'
       , testProperty "strict foldl'" pStrictFoldl'

containers/src/Data/IntMap/Internal.hs

@@ -2582,13 +2582,10 @@ filter p 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 = 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)
+filterKeys predicate = filterWithKey (\k _ -> predicate k)
 
 -- | \(O(n)\). Filter all keys\/values that satisfy some predicate.
 --