diff --git a/containers-tests/tests/intmap-strictness.hs b/containers-tests/tests/intmap-strictness.hs index 44f4d1260..0e42c43aa 100644 --- a/containers-tests/tests/intmap-strictness.hs +++ b/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' diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs index 1fd741c77..c9be1b56f 100644 --- a/containers/src/Data/IntMap/Internal.hs +++ b/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. --