diff --git a/Data/IntMap/Internal.hs b/Data/IntMap/Internal.hs index 9b30673cc..4132545ca 100644 --- a/Data/IntMap/Internal.hs +++ b/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 diff --git a/benchmarks/IntMap.hs b/benchmarks/IntMap.hs index 8fabda4bc..43f7e0372 100644 --- a/benchmarks/IntMap.hs +++ b/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 diff --git a/tests/intmap-properties.hs b/tests/intmap-properties.hs index 18c55e675..4463e43b8 100644 --- a/tests/intmap-properties.hs +++ b/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)