Add RULES for Data.IntMap.alterF

Data/IntMap/Internal.hs

@@ -286,14 +286,19 @@ module Data.IntMap.Internal (
 
 #if MIN_VERSION_base(4,8,0)
 import Data.Functor.Identity (Identity (..))
+#if MIN_VERSION_base(4,9,0)
+import Control.Applicative (liftA2, Const(..))
+#else
 import Control.Applicative (liftA2)
+#endif
 #else
 import Control.Applicative (Applicative(pure, (<*>)), (<$>), liftA2)
 import Data.Monoid (Monoid(..))
 import Data.Traversable (Traversable(traverse))
 import Data.Word (Word)
 #endif
 #if MIN_VERSION_base(4,9,0)
+import Data.Coerce (coerce)
 import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid)
 import Data.Functor.Classes
 #endif
@@ -980,6 +985,22 @@ 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
+-- TODO(m-renaud): Figure out if this should be marked INLINE or NOINLINE.
+-- It needs to be one or the other or else the specialization rule may not fire.
+{-# NOINLINE [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
+
+#if MIN_VERSION_base(4,9,0)
+{-# RULES
+"Const specialize alterF" forall (f :: Maybe a -> Const x (Maybe a)) k m.
+  alterF f k m = coerce . f $ lookup k m
+  #-}
+#endif
 
 {--------------------------------------------------------------------
   Union

benchmarks/IntMap.hs

@@ -1,9 +1,16 @@
 {-# 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.Functor.Identity (Identity (..))
+#endif
 import Data.List (foldl')
 import qualified Data.IntMap as M
 import qualified Data.IntMap.Strict as MS
@@ -35,6 +42,13 @@ main = do
         , bench "update" $ whnf (upd keys) m
         , bench "updateLookupWithKey" $ whnf (upd' keys) m
         , bench "alter"  $ whnf (alt keys) m
+        , bench "alterF" $ whnf (altFList keys) m
+#if MIN_VERSION_base(4,8,0)
+        , bench "alterFIdentity" $ whnf (altFIdentity keys) m
+#endif
+#if MIN_VERSION_base(4,9,0)
+        , bench "alterFConst" $ 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 +104,22 @@ 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
 
+altFList :: [Int] -> M.IntMap Int -> M.IntMap Int
+altFList xs m = foldl' (\m k -> head $ M.alterF (pure . id) k m) m xs
+
+#if MIN_VERSION_base(4,8,0)
+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

@@ -7,6 +7,12 @@ import Data.IntMap.Lazy as Data.IntMap hiding (showTree)
 #endif
 import Data.IntMap.Internal.Debug (showTree)
 
+#if MIN_VERSION_base(4,9,0)
+import Control.Applicative (Const(..))
+#endif
+#if MIN_VERSION_base(4,8,0)
+import Data.Functor.Identity (Identity (..))
+#endif
 import Data.Monoid
 import Data.Maybe hiding (mapMaybe)
 import qualified Data.Maybe as Maybe (mapMaybe)
@@ -56,6 +62,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
@@ -143,6 +150,12 @@ 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
+#if MIN_VERSION_base(4,9,0)
+             , testProperty "alterF_Const"         prop_alterF_ConstRules
+#endif
              , testProperty "index"                prop_index
              , testProperty "index_lookup"         prop_index_lookup
              , testProperty "null"                 prop_null
@@ -399,9 +412,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
 
@@ -922,6 +976,52 @@ prop_alter t k = 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 f (TestIdentity a) = TestIdentity (f a)
+
+prop_alterF_IdentityRules :: UMap -> Int -> Bool
+prop_alterF_IdentityRules t k =
+    runIdentity tIdentity == runTestIdentity tTestIdentity
+  where
+    tIdentity = alterF fIdentity k t
+    fIdentity Nothing = Identity (Just ())
+    fIdentity (Just ()) = Identity Nothing
+
+    tTestIdentity = alterF fTest k t
+    fTest Nothing   = TestIdentity (Just ())
+    fTest (Just ()) = TestIdentity (Nothing)
+#endif
+
+#if MIN_VERSION_base(4,9,0)
+-- Verify that the rewrite rules for Const give the same result
+-- as the non-rewritten version. We use a custom TestConst that
+-- will not fire the rewrite rules to compare against.
+
+newtype TestConst a b = TestConst { getTestConst :: a }
+
+instance Functor (TestConst a) where
+    fmap _ (TestConst a) = TestConst a
+
+prop_alterF_ConstRules :: UMap -> Int -> Bool
+prop_alterF_ConstRules t k =
+    getConst tConst == getTestConst tTestConst
+  where
+    tConst = alterF fConst k t
+    fConst Nothing = Const False
+    fConst (Just ()) = Const True
+
+    tTestConst = alterF fTest k t
+    fTest Nothing   = TestConst False
+    fTest (Just ()) = TestConst True
+#endif
+
 ------------------------------------------------------------------------
 -- Compare against the list model (after nub on keys)