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
@@ -12,6 +20,8 @@ import Prelude hiding (lookup)
 
 main = do
     let m = M.fromAscList elems :: M.IntMap Int
+        m_even = M.fromAscList elems_even :: M.IntMap Int
+        m_odd = M.fromAscList elems_odd :: M.IntMap Int
     evaluate $ rnf [m]
     defaultMain
         [ bench "lookup" $ whnf (lookup keys) m
@@ -35,6 +45,26 @@ main = do
         , bench "update" $ whnf (upd keys) m
         , bench "updateLookupWithKey" $ whnf (upd' keys) m
         , bench "alter"  $ whnf (alt keys) m
+        , bench "alterF lookup absent" $ whnf (atLookup evens) m_odd
+        , bench "alterF lookup present" $ whnf (atLookup evens) m_even
+        , bench "alterF no rules lookup absent" $ whnf (atLookupNoRules evens) m_odd
+        , bench "alterF no rules lookup present" $ whnf (atLookupNoRules evens) m_even
+        , bench "alterF insert absent" $ whnf (atIns elems_even) m_odd
+        , bench "alterF insert present" $ whnf (atIns elems_even) m_even
+        , bench "alterF no rules insert absent" $ whnf (atInsNoRules elems_even) m_odd
+        , bench "alterF no rules insert present" $ whnf (atInsNoRules elems_even) m_even
+        , bench "alterF delete absent" $ whnf (atDel evens) m_odd
+        , bench "alterF delete present" $ whnf (atDel evens) m
+        , bench "alterF no rules delete absent" $ whnf (atDelNoRules evens) m_odd
+        , bench "alterF no rules delete present" $ whnf (atDelNoRules evens) m
+        , bench "alterF alter absent" $ whnf (atAlt id evens) m_odd
+        , bench "alterF alter insert" $ whnf (atAlt (const (Just 1)) evens) m_odd
+        , bench "alterF alter update" $ whnf (atAlt id evens) m_even
+        , bench "alterF alter delete" $ whnf (atAlt (const Nothing) evens) m
+        , bench "alterF no rules alter absent" $ whnf (atAltNoRules id evens) m_odd
+        , bench "alterF no rules alter insert" $ whnf (atAltNoRules (const (Just 1)) evens) m_odd
+        , bench "alterF no rules alter update" $ whnf (atAltNoRules id evens) m_even
+        , bench "alterF no rules alter delete" $ whnf (atAltNoRules (const Nothing) evens) m
         , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
         , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
         , bench "fromList" $ whnf M.fromList elems
@@ -43,7 +73,11 @@ main = do
         ]
   where
     elems = zip keys values
+    elems_even = zip evens evens
+    elems_odd = zip odds odds
     keys = [1..2^12]
+    evens = [2,4..2^12]
+    odds = [1,3..2^12]
     values = [1..2^12]
     sum k v1 v2 = k + v1 + v2
     consPair k v xs = (k, v) : xs
@@ -90,6 +124,48 @@ 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
 
+newtype TestIdentity a = TestIdentity { runTestIdentity :: a }
+
+instance Functor TestIdentity where
+#if MIN_VERSION_base(4,8,0)
+  fmap = coerce
+#else
+  fmap f (Ident a) = Ident (f a)
+#endif
+
+newtype TestConst a b = TestConst { getTestConst :: a }
+instance Functor (TestConst a) where
+  fmap _ (TestConst a) = TestConst a
+
+atLookup :: [Int] -> M.IntMap Int -> Int
+atLookup xs m = foldl' (\n k -> fromMaybe n (getConst (M.alterF Const k m))) 0 xs
+
+atLookupNoRules :: [Int] -> M.IntMap Int -> Int
+atLookupNoRules xs m =
+  foldl' (\n k -> fromMaybe n (getTestConst (M.alterF TestConst k m))) 0 xs
+
+atIns :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int
+atIns xs m =
+  foldl' (\m (k, v) -> runIdentity (M.alterF (\_ -> Identity (Just v)) k m)) m xs
+
+atInsNoRules :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int
+atInsNoRules xs m =
+  foldl' (\m (k, v) -> runTestIdentity (M.alterF (\_ -> TestIdentity (Just v)) k m)) m xs
+
+atDel :: [Int] -> M.IntMap Int -> M.IntMap Int
+atDel xs m = foldl' (\m k -> runIdentity (M.alterF (\_ -> Identity Nothing) k m)) m xs
+
+atDelNoRules :: [Int] -> M.IntMap Int -> M.IntMap Int
+atDelNoRules xs m =
+  foldl' (\m k -> runTestIdentity (M.alterF (\_ -> TestIdentity Nothing) k m)) m xs
+
+atAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.IntMap Int -> M.IntMap Int
+atAlt f xs m = foldl' (\m k -> runIdentity (M.alterF (Identity . f) k m)) m xs
+
+atAltNoRules :: (Maybe Int -> Maybe Int) -> [Int] -> M.IntMap Int -> M.IntMap Int
+atAltNoRules f xs m =
+  foldl' (\m k -> runTestIdentity (M.alterF (TestIdentity . f) k m)) m xs
+
 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)