Skip to content

Commit 9d90611

Browse files
authored
Add filterKeys for Map and IntMap (#972)
1 parent b24068b commit 9d90611

File tree

10 files changed

+81
-7
lines changed

10 files changed

+81
-7
lines changed

containers-tests/tests/intmap-properties.hs

+24-4
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import Data.Foldable (foldMap)
2323
import Data.Function
2424
import Data.Traversable (Traversable(traverse), foldMapDefault)
2525
import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl')
26-
import qualified Prelude (map)
26+
import qualified Prelude (map, filter)
2727

2828
import Data.List (nub,sort)
2929
import qualified Data.List as List
@@ -103,7 +103,8 @@ main = defaultMain $ testGroup "intmap-properties"
103103
, testCase "fromAscListWithKey" test_fromAscListWithKey
104104
, testCase "fromDistinctAscList" test_fromDistinctAscList
105105
, testCase "filter" test_filter
106-
, testCase "filterWithKey" test_filteWithKey
106+
, testCase "filterKeys" test_filterKeys
107+
, testCase "filterWithKey" test_filterWithKey
107108
, testCase "partition" test_partition
108109
, testCase "partitionWithKey" test_partitionWithKey
109110
, testCase "mapMaybe" test_mapMaybe
@@ -182,6 +183,8 @@ main = defaultMain $ testGroup "intmap-properties"
182183
, testProperty "deleteMin" prop_deleteMinModel
183184
, testProperty "deleteMax" prop_deleteMaxModel
184185
, testProperty "filter" prop_filter
186+
, testProperty "filterKeys" prop_filterKeys
187+
, testProperty "filterWithKey" prop_filterWithKey
185188
, testProperty "partition" prop_partition
186189
, testProperty "takeWhileAntitone" prop_takeWhileAntitone
187190
, testProperty "dropWhileAntitone" prop_dropWhileAntitone
@@ -867,8 +870,13 @@ test_filter = do
867870
filter (> "x") (fromList [(5,"a"), (-3,"b")]) @?= empty
868871
filter (< "a") (fromList [(5,"a"), (-3,"b")]) @?= empty
869872

870-
test_filteWithKey :: Assertion
871-
test_filteWithKey = do
873+
test_filterKeys :: Assertion
874+
test_filterKeys = do
875+
filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
876+
filterKeys (> 4) (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "a"
877+
878+
test_filterWithKey :: Assertion
879+
test_filterWithKey = do
872880
filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
873881
filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "a"
874882

@@ -1506,6 +1514,18 @@ prop_filter p ys = length ys > 0 ==>
15061514
in valid m .&&.
15071515
m === fromList (List.filter (apply p . snd) xs)
15081516

1517+
prop_filterKeys :: Fun Int Bool -> IMap -> Property
1518+
prop_filterKeys fun m =
1519+
valid m' .&&. toList m' === Prelude.filter (apply fun . fst) (toList m)
1520+
where
1521+
m' = filterKeys (apply fun) m
1522+
1523+
prop_filterWithKey :: Fun (Int, Int) Bool -> IMap -> Property
1524+
prop_filterWithKey fun m =
1525+
valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
1526+
where
1527+
m' = filterWithKey (applyFun2 fun) m
1528+
15091529
prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
15101530
prop_partition p ys = length ys > 0 ==>
15111531
let xs = List.nubBy ((==) `on` fst) ys

containers-tests/tests/map-properties.hs

+21-3
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,8 @@ main = defaultMain $ testGroup "map-properties"
118118
, testCase "fromDistinctAscList" test_fromDistinctAscList
119119
, testCase "fromDistinctDescList" test_fromDistinctDescList
120120
, testCase "filter" test_filter
121-
, testCase "filterWithKey" test_filteWithKey
121+
, testCase "filterKeys" test_filterKeys
122+
, testCase "filterWithKey" test_filterWithKey
122123
, testCase "partition" test_partition
123124
, testCase "partitionWithKey" test_partitionWithKey
124125
, testCase "mapMaybe" test_mapMaybe
@@ -223,6 +224,8 @@ main = defaultMain $ testGroup "map-properties"
223224
, testProperty "deleteMin" prop_deleteMinModel
224225
, testProperty "deleteMax" prop_deleteMaxModel
225226
, testProperty "filter" prop_filter
227+
, testProperty "filterKeys" prop_filterKeys
228+
, testProperty "filterWithKey" prop_filterWithKey
226229
, testProperty "partition" prop_partition
227230
, testProperty "map" prop_map
228231
, testProperty "fmap" prop_fmap
@@ -796,8 +799,11 @@ test_filter = do
796799
filter (> "x") (fromList [(5,"a"), (3,"b")]) @?= empty
797800
filter (< "a") (fromList [(5,"a"), (3,"b")]) @?= empty
798801

799-
test_filteWithKey :: Assertion
800-
test_filteWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
802+
test_filterKeys :: Assertion
803+
test_filterKeys = filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
804+
805+
test_filterWithKey :: Assertion
806+
test_filterWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a"
801807

802808
test_partition :: Assertion
803809
test_partition = do
@@ -1463,6 +1469,18 @@ prop_filter p ys = length ys > 0 ==>
14631469
m = fromList xs
14641470
in filter (apply p) m == fromList (List.filter (apply p . snd) xs)
14651471

1472+
prop_filterKeys :: Fun Int Bool -> IMap -> Property
1473+
prop_filterKeys fun m =
1474+
valid m' .&&. toList m' === Prelude.filter (apply fun . fst) (toList m)
1475+
where
1476+
m' = filterKeys (apply fun) m
1477+
1478+
prop_filterWithKey :: Fun (Int, Int) Bool -> IMap -> Property
1479+
prop_filterWithKey fun m =
1480+
valid m' .&&. toList m' === Prelude.filter (apply fun) (toList m)
1481+
where
1482+
m' = filterWithKey (apply2 fun) m
1483+
14661484
prop_take :: Int -> Map Int Int -> Property
14671485
prop_take n xs = valid taken .&&.
14681486
taken === fromDistinctAscList (List.take n (toList xs))

containers/src/Data/IntMap/Internal.hs

+14
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,7 @@ module Data.IntMap.Internal (
218218

219219
-- * Filter
220220
, filter
221+
, filterKeys
221222
, filterWithKey
222223
, restrictKeys
223224
, withoutKeys
@@ -2638,6 +2639,19 @@ filter :: (a -> Bool) -> IntMap a -> IntMap a
26382639
filter p m
26392640
= filterWithKey (\_ x -> p x) m
26402641

2642+
-- | \(O(n)\). Filter all keys that satisfy some predicate.
2643+
--
2644+
-- @
2645+
-- filterKeys p = 'filterWithKey' (\\k _ -> p k)
2646+
-- @
2647+
--
2648+
-- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2649+
--
2650+
-- @since FIXME
2651+
2652+
filterKeys :: (Key -> Bool) -> IntMap a -> IntMap a
2653+
filterKeys predicate = filterWithKey (\k _ -> predicate k)
2654+
26412655
-- | \(O(n)\). Filter all keys\/values that satisfy some predicate.
26422656
--
26432657
-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"

containers/src/Data/IntMap/Lazy.hs

+1
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,7 @@ module Data.IntMap.Lazy (
198198

199199
-- * Filter
200200
, IM.filter
201+
, filterKeys
201202
, filterWithKey
202203
, restrictKeys
203204
, withoutKeys

containers/src/Data/IntMap/Strict.hs

+1
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ module Data.IntMap.Strict (
216216

217217
-- * Filter
218218
, filter
219+
, filterKeys
219220
, filterWithKey
220221
, restrictKeys
221222
, withoutKeys

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

+2
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ module Data.IntMap.Strict.Internal (
214214

215215
-- * Filter
216216
, filter
217+
, filterKeys
217218
, filterWithKey
218219
, restrictKeys
219220
, withoutKeys
@@ -279,6 +280,7 @@ import Data.IntMap.Internal
279280
, empty
280281
, assocs
281282
, filter
283+
, filterKeys
282284
, filterWithKey
283285
, findMin
284286
, findMax

containers/src/Data/Map/Internal.hs

+14
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,7 @@ module Data.Map.Internal (
291291

292292
-- * Filter
293293
, filter
294+
, filterKeys
294295
, filterWithKey
295296

296297
, takeWhileAntitone
@@ -2944,6 +2945,19 @@ filter :: (a -> Bool) -> Map k a -> Map k a
29442945
filter p m
29452946
= filterWithKey (\_ x -> p x) m
29462947

2948+
-- | \(O(n)\). Filter all keys that satisfy the predicate.
2949+
--
2950+
-- @
2951+
-- filterKeys p = 'filterWithKey' (\\k _ -> p k)
2952+
-- @
2953+
--
2954+
-- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
2955+
--
2956+
-- @since FIXME
2957+
2958+
filterKeys :: (k -> Bool) -> Map k a -> Map k a
2959+
filterKeys p m = filterWithKey (\k _ -> p k) m
2960+
29472961
-- | \(O(n)\). Filter all keys\/values that satisfy the predicate.
29482962
--
29492963
-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"

containers/src/Data/Map/Lazy.hs

+1
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,7 @@ module Data.Map.Lazy (
228228

229229
-- * Filter
230230
, filter
231+
, filterKeys
231232
, filterWithKey
232233
, restrictKeys
233234
, withoutKeys

containers/src/Data/Map/Strict.hs

+1
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@ module Data.Map.Strict
243243

244244
-- * Filter
245245
, filter
246+
, filterKeys
246247
, filterWithKey
247248
, restrictKeys
248249
, withoutKeys

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

+2
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ module Data.Map.Strict.Internal
253253

254254
-- * Filter
255255
, filter
256+
, filterKeys
256257
, filterWithKey
257258
, restrictKeys
258259
, withoutKeys
@@ -359,6 +360,7 @@ import Data.Map.Internal
359360
, drop
360361
, dropWhileAntitone
361362
, filter
363+
, filterKeys
362364
, filterWithKey
363365
, findIndex
364366
, findMax

0 commit comments

Comments
 (0)