Add RULES for Data.IntMap.alterF

Data/IntMap/Internal.hs

@@ -991,6 +991,13 @@ alterF f k m = (<$> f mv) $ \fres ->
     Nothing -> maybe m (const (delete k m)) mv
     Just v' -> insert k v' m
   where mv = lookup k m
+{-# INLINABLE [1] alterF #-}
+#if MIN_VERSION_base(4,8,0)
+{-# RULES
+"Identity specialize alterF" forall (f :: Maybe a -> Identity (Maybe a)) k m.
+  alterF f k m = Identity $ alter (coerce f) k m
+  #-}
+#endif
 
 {--------------------------------------------------------------------
   Union

benchmarks/IntMap.hs

@@ -1,9 +1,17 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
 module Main where
 
+#if MIN_VERSION_base(4,9,0)
+import Control.Applicative (Const(..))
+#endif
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
 import Criterion.Main (bench, defaultMain, whnf)
+#if MIN_VERSION_base(4,8,0)
+import Data.Coerce (coerce)
+import Data.Functor.Identity (Identity (..))
+#endif
 import Data.List (foldl')
 import qualified Data.IntMap as M
 import qualified Data.IntMap.Strict as MS
@@ -35,6 +43,13 @@ main = do
         , bench "update" $ whnf (upd keys) m
         , bench "updateLookupWithKey" $ whnf (upd' keys) m
         , bench "alter"  $ whnf (alt keys) m
+#if MIN_VERSION_base(4,8,0)
+        , bench "alterF_IdentityNoRewrite" $ whnf (altFTestIdentity keys) m
+        , bench "alterF_IdentityRewrite" $ whnf (altFIdentity keys) m
+#endif
+#if MIN_VERSION_base(4,9,0)
+        , bench "alterF_Const" $ whnf (altFConst keys) m
+#endif
         , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
         , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
         , bench "fromList" $ whnf M.fromList elems
@@ -90,6 +105,28 @@ upd' xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> Just a) k m) m
 alt :: [Int] -> M.IntMap Int -> M.IntMap Int
 alt xs m = foldl' (\m k -> M.alter id k m) m xs
 
+#if MIN_VERSION_base(4,8,0)
+newtype TestIdentity a = TestIdentity { runTestIdentity :: a }
+
+instance Functor TestIdentity where
+    fmap = coerce
+
+altFTestIdentity :: [Int] -> M.IntMap Int -> M.IntMap Int
+altFTestIdentity xs m =
+    foldl' (\m k -> runTestIdentity $ M.alterF (TestIdentity . id) k m) m xs
+
+altFIdentity :: [Int] -> M.IntMap Int -> M.IntMap Int
+altFIdentity xs m = foldl' (\m k -> runIdentity $ M.alterF (pure . id) k m) m xs
+#endif
+
+#if MIN_VERSION_base(4,9,0)
+altFConst :: [Int] -> M.IntMap Int -> M.IntMap Int
+altFConst xs m =
+    foldl' (\m k -> getConst $ M.alterF (const (Const m) . id) k m) m xs
+#endif
+
+
+
 maybeDel :: Int -> Maybe Int
 maybeDel n | n `mod` 3 == 0 = Nothing
            | otherwise      = Just n

tests/intmap-properties.hs

@@ -8,6 +8,13 @@ import Data.IntMap.Lazy as Data.IntMap hiding (showTree)
 import Data.IntMap.Internal.Debug (showTree)
 import IntMapValidity (valid)
 
+#if MIN_VERSION_base(4,9,0)
+import Control.Applicative (Const(..))
+#endif
+#if MIN_VERSION_base(4,8,0)
+import Data.Coerce (coerce)
+import Data.Functor.Identity (Identity (..))
+#endif
 import Data.Monoid
 import Data.Maybe hiding (mapMaybe)
 import qualified Data.Maybe as Maybe (mapMaybe)
@@ -57,6 +64,7 @@ main = defaultMain
              , testCase "updateWithKey" test_updateWithKey
              , testCase "updateLookupWithKey" test_updateLookupWithKey
              , testCase "alter" test_alter
+             , testCase "alterF" test_alterF
              , testCase "union" test_union
              , testCase "mappend" test_mappend
              , testCase "unionWith" test_unionWith
@@ -146,6 +154,9 @@ main = defaultMain
              , testProperty "toAscList+toDescList" prop_ascDescList
              , testProperty "fromList"             prop_fromList
              , testProperty "alter"                prop_alter
+#if MIN_VERSION_base(4,8,0)
+             , testProperty "alterF_Identity"      prop_alterF_IdentityRules
+#endif
              , testProperty "index"                prop_index
              , testProperty "index_lookup"         prop_index_lookup
              , testProperty "null"                 prop_null
@@ -402,9 +413,50 @@ test_alter = do
     alter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")]
     alter g 5 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "c")]
   where
+    f, g :: Maybe String -> Maybe String
     f _ = Nothing
     g _ = Just "c"
 
+test_alterF :: Assertion
+test_alterF = do
+    let m = fromList [(5,"a"), (3,"b")]
+    -- List applicative
+    alterF fList 7 m @?= [fromList [(3, "b"), (5, "a")]]
+    alterF fList 5 m @?= [singleton 3 "b"]
+    alterF gList 7 m @?= [fromList [(3, "b"), (5, "a"), (7, "c")]]
+    alterF gList 5 m @?= [fromList [(3, "b"), (5, "c")]]
+#if MIN_VERSION_base(4,8,0)
+    -- Identity applicative
+    alterF fIdentity 7 m @?= Identity (fromList [(3, "b"), (5, "a")])
+    alterF fIdentity 5 m @?= Identity (singleton 3 "b")
+    alterF gIdentity 7 m @?= Identity (fromList [(3, "b"), (5, "a"), (7, "c")])
+    alterF gIdentity 5 m @?= Identity (fromList [(3, "b"), (5, "c")])
+#endif
+#if MIN_VERSION_base(4,9,0)
+    -- Const applicative
+    alterF fConst 7 m @?= Const False
+    alterF fConst 5 m @?= Const False
+    alterF gConst 7 m @?= Const True
+    alterF gConst 5 m @?= Const True
+#endif
+  where
+    fList, gList :: Maybe String -> [Maybe String]
+    fList _ = [Nothing]
+    gList _ = [Just "c"]
+
+#if MIN_VERSION_base(4,8,0)
+    fIdentity, gIdentity :: Maybe String -> Identity (Maybe String)
+    fIdentity _ = Identity Nothing
+    gIdentity _ = Identity (Just "c")
+#endif
+
+#if MIN_VERSION_base(4,9,0)
+    fConst, gConst :: Maybe String -> Const Bool (Maybe String)
+    fConst _ = Const False
+    gConst _ = Const True
+#endif
+
+
 ----------------------------------------------------------------
 -- Combine
 
@@ -963,6 +1015,31 @@ prop_alter t k = valid t' .&&. case lookup k t of
     f Nothing   = Just ()
     f (Just ()) = Nothing
 
+#if MIN_VERSION_base(4,8,0)
+-- Verify that the rewrite rules for Identity give the same result as the
+-- non-rewritten version. We use our own TestIdentity functor to compare
+-- against.
+
+newtype TestIdentity a = TestIdentity { runTestIdentity :: a }
+
+instance Functor TestIdentity where
+    fmap = coerce
+
+prop_alterF_IdentityRules :: UMap -> Int -> Property
+prop_alterF_IdentityRules t k =
+    valid tIdentity .&&.
+    valid tTestIdentity .&&.
+    tIdentity == tTestIdentity
+  where
+    tIdentity = runIdentity $ alterF fIdentity k t
+    fIdentity Nothing = Identity (Just ())
+    fIdentity (Just ()) = Identity Nothing
+
+    tTestIdentity = runTestIdentity $ alterF fTest k t
+    fTest Nothing   = TestIdentity (Just ())
+    fTest (Just ()) = TestIdentity (Nothing)
+#endif
+
 ------------------------------------------------------------------------
 -- Compare against the list model (after nub on keys)