Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add filterKeys for Map and IntMap #972

Merged
merged 11 commits into from
Nov 23, 2024
35 changes: 31 additions & 4 deletions containers-tests/tests/intmap-properties.hs
flip111 marked this conversation as resolved.
Show resolved Hide resolved
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"

Expand Down Expand Up @@ -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
flip111 marked this conversation as resolved.
Show resolved Hide resolved

prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
prop_partition p ys = length ys > 0 ==>
let xs = List.nubBy ((==) `on` fst) ys
Expand Down
10 changes: 7 additions & 3 deletions containers-tests/tests/map-properties.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
meooow25 marked this conversation as resolved.
Show resolved Hide resolved

test_filterWithKey :: Assertion
test_filterWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"

test_partition :: Assertion
test_partition = do
Expand Down
9 changes: 9 additions & 0 deletions containers/src/Data/IntMap/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,7 @@ module Data.IntMap.Internal (

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down Expand Up @@ -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"
meooow25 marked this conversation as resolved.
Show resolved Hide resolved
-- > filterKeys (> 4) == filterWithKey (\k _ -> k > 4)
meooow25 marked this conversation as resolved.
Show resolved Hide resolved

filterKeys :: (Key -> Bool) -> IntMap a -> IntMap a
filterKeys predicate = filterWithKey (\k _ -> predicate k)
flip111 marked this conversation as resolved.
Show resolved Hide resolved

-- | \(O(n)\). Filter all keys\/values that satisfy some predicate.
--
-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
Expand Down
1 change: 1 addition & 0 deletions containers/src/Data/IntMap/Lazy.hs
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,7 @@ module Data.IntMap.Lazy (

-- * Filter
, IM.filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down
1 change: 1 addition & 0 deletions containers/src/Data/IntMap/Strict.hs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ module Data.IntMap.Strict (

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down
2 changes: 2 additions & 0 deletions containers/src/Data/IntMap/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ module Data.IntMap.Strict.Internal (

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down Expand Up @@ -286,6 +287,7 @@ import Data.IntMap.Internal
, empty
, assocs
, filter
, filterKeys
, filterWithKey
, findMin
, findMax
Expand Down
9 changes: 9 additions & 0 deletions containers/src/Data/Map/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,7 @@ module Data.Map.Internal (

-- * Filter
, filter
, filterKeys
, filterWithKey

, takeWhileAntitone
Expand Down Expand Up @@ -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

meooow25 marked this conversation as resolved.
Show resolved Hide resolved
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"
Expand Down
1 change: 1 addition & 0 deletions containers/src/Data/Map/Lazy.hs
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ module Data.Map.Lazy (

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down
1 change: 1 addition & 0 deletions containers/src/Data/Map/Strict.hs
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,7 @@ module Data.Map.Strict

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down
2 changes: 2 additions & 0 deletions containers/src/Data/Map/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,7 @@ module Data.Map.Strict.Internal

-- * Filter
, filter
, filterKeys
, filterWithKey
, restrictKeys
, withoutKeys
Expand Down Expand Up @@ -359,6 +360,7 @@ import Data.Map.Internal
, drop
, dropWhileAntitone
, filter
, filterKeys
, filterWithKey
, findIndex
, findMax
Expand Down
Loading