From 09ebe0f45a7873920581dc318cdc49cf14c9bd0f Mon Sep 17 00:00:00 2001
From: flip111 <flip101@gmail.com>
Date: Sun, 17 Nov 2024 23:42:10 +0100
Subject: [PATCH] Changes after review 2

---
 containers-tests/tests/intmap-properties.hs | 23 ++++++++++++++++++++-
 containers-tests/tests/intmap-strictness.hs | 19 -----------------
 containers-tests/tests/map-properties.hs    |  1 +
 containers/src/Data/Map/Internal.hs         |  1 +
 4 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/containers-tests/tests/intmap-properties.hs b/containers-tests/tests/intmap-properties.hs
index b4c5b4ae1..fbe2de47c 100644
--- a/containers-tests/tests/intmap-properties.hs
+++ b/containers-tests/tests/intmap-properties.hs
@@ -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
@@ -180,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
@@ -1470,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
+
 prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
 prop_partition p ys = length ys > 0 ==>
   let xs = List.nubBy ((==) `on` fst) ys
diff --git a/containers-tests/tests/intmap-strictness.hs b/containers-tests/tests/intmap-strictness.hs
index 0e42c43aa..44f4d1260 100644
--- a/containers-tests/tests/intmap-strictness.hs
+++ b/containers-tests/tests/intmap-strictness.hs
@@ -87,24 +87,6 @@ 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
 --
@@ -217,7 +199,6 @@ 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-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs
index ec075b1b3..e2f052a42 100644
--- a/containers-tests/tests/map-properties.hs
+++ b/containers-tests/tests/map-properties.hs
@@ -115,6 +115,7 @@ main = defaultMain $ testGroup "map-properties"
          , testCase "fromDistinctAscList" test_fromDistinctAscList
          , testCase "fromDistinctDescList" test_fromDistinctDescList
          , testCase "filter" test_filter
+         , testCase "filterKeys" test_filterKeys
          , testCase "filterWithKey" test_filterWithKey
          , testCase "partition" test_partition
          , testCase "partitionWithKey" test_partitionWithKey
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index b8d047a08..3c01337bf 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/containers/src/Data/Map/Internal.hs
@@ -2861,6 +2861,7 @@ filter p m
 -- | \(O(n)\). Filter all keys that satisfy the predicate.
 --
 -- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
+-- @since FIXME
 
 filterKeys :: (k -> Bool) -> Map k a -> Map k a
 filterKeys p m = filterWithKey (\k _ -> p k) m