Add rewrite rule for alterF with pairs
authorDavid Feuer <David.Feuer@gmail.com>
Sun, 22 May 2016 22:27:57 +0000 (18:27 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Sun, 22 May 2016 22:41:13 +0000 (18:41 -0400)
Also fix alterF documentation formatting.

Data/Map/Base.hs
Data/Map/Strict.hs
benchmarks/Map.hs

index 92ecda0..2e0bead 100644 (file)
@@ -267,6 +267,9 @@ module Data.Map.Base (
     -- Used by the strict version
     , AreWeStrict (..)
     , atKeyImpl
+#if __GLASGOW_HASKELL__
+    , atKeyWithLookup
+#endif
 #if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0)
     , atKeyPlain
 #endif
@@ -973,6 +976,7 @@ data AreWeStrict = Strict | Lazy
 -- ('lookup' k m)@.
 --
 -- Example:
+--
 -- @
 -- interactiveAlter :: Int -> Map Int String -> IO (Map Int String)
 -- interactiveAlter k m = alterF f k m where
@@ -995,13 +999,12 @@ data AreWeStrict = Strict | Lazy
 --
 -- Note on rewrite rules:
 --
--- This module includes GHC rewrite rules to optimize 'alterF' for
--- the 'Const' and 'Identity' functors. In general, these rules
--- improve performance. The sole exception is that when using
--- 'Identity', deleting a key that is already absent takes longer
--- than it would without the rules. If you expect this to occur
--- a very large fraction of the time, you might consider using a
--- private copy of the 'Identity' type.
+-- This module includes GHC rewrite rules to optimize 'alterF' for the 'Const',
+-- 'Identity', and @(,) b@ functors. In general, these rules improve
+-- performance. The main exception is that when using 'Identity', deleting a
+-- key that is already absent takes longer than it would without the rules. If
+-- you expect this to occur a very large fraction of the time, you might
+-- consider using a private copy of the 'Identity' type.
 --
 -- Note: 'alterF' is a flipped version of the 'at' combinator from
 -- 'Control.Lens.At'.
@@ -1017,10 +1020,14 @@ alterF f k m = atKeyImpl Lazy k f m
 {-# INLINABLE [2] alterF #-}
 
 -- We can save a little time by recognizing the special case of
--- `Control.Applicative.Const` and just doing a lookup.
+-- `Control.Applicative.Const` and just doing a lookup. Similarly,
+-- we recognize the special case of `(,) b` which, like `Identity`,
+-- only needs to go down and up once.
 {-# RULES
 "alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m
+"alterF/Pair" forall k (f :: Maybe a -> (b, Maybe a)) . alterF f k = atKeyPair k f
  #-}
+
 #if MIN_VERSION_base(4,8,0)
 -- base 4.8 and above include Data.Functor.Identity, so we can
 -- save a pretty decent amount of time by handling it specially.
@@ -1195,6 +1202,46 @@ atKeyPlain strict k0 f0 t = case go k0 f0 t of
 data Altered k a = AltSmaller !(Map k a) | AltBigger !(Map k a) | AltAdj !(Map k a) | AltSame
 #endif
 
+#if __GLASGOW_HASKELL__
+atKeyPair :: Ord k => k -> (Maybe a -> (b, Maybe a)) -> Map k a -> (b, Map k a)
+atKeyPair k f t = atKeyWithLookup Lazy k f t
+{-# INLINABLE atKeyPair #-}
+
+atKeyWithLookup :: Ord k => AreWeStrict -> k -> (Maybe a -> (b, Maybe a)) -> Map k a -> (b, Map k a)
+atKeyWithLookup strict k0 f0 t = case go k0 f0 t of
+    AltSmallerLook v t' -> (v, t')
+    AltBiggerLook v t' -> (v, t')
+    AltAdjLook v t' -> (v, t')
+    AltSameLook v -> (v, t)
+  where
+    go :: Ord k => k -> (Maybe a -> (b, Maybe a)) -> Map k a -> AlteredLookup k b a
+    go !k f Tip = case f Nothing of
+                   (b, Nothing) -> AltSameLook b
+                   (b, Just x)  -> case strict of
+                     Lazy -> AltBiggerLook b (singleton k x)
+                     Strict -> x `seq` (AltBiggerLook b $ singleton k x)
+
+    go k f (Bin sx kx x l r) = case compare k kx of
+                   LT -> case go k f l of
+                           AltSmallerLook b l' -> AltSmallerLook b $ balanceR kx x l' r
+                           AltBiggerLook b l' -> AltBiggerLook b $ balanceL kx x l' r
+                           AltAdjLook b l' -> AltAdjLook b $ Bin sx kx x l' r
+                           s@AltSameLook{} -> s
+                   GT -> case go k f r of
+                           AltSmallerLook b r' -> AltSmallerLook b $ balanceL kx x l r'
+                           AltBiggerLook b r' -> AltBiggerLook b $ balanceR kx x l r'
+                           AltAdjLook b r' -> AltAdjLook b $ Bin sx kx x l r'
+                           s@AltSameLook{} -> s
+                   EQ -> case f (Just x) of
+                           (b, Just x') -> case strict of
+                             Lazy -> AltAdjLook b $ Bin sx kx x' l r
+                             Strict -> x' `seq` (AltAdjLook b $ Bin sx kx x' l r)
+                           (b, Nothing) -> AltSmallerLook b $ glue l r
+{-# INLINE atKeyWithLookup #-}
+
+data AlteredLookup k b a = AltSmallerLook b !(Map k a) | AltBiggerLook b !(Map k a) | AltAdjLook b !(Map k a) | AltSameLook b
+#endif
+
 #if DEFINE_ALTERF_FALLBACK
 -- When the map is too large to use a bit queue, we fall back to
 -- this much slower version which uses a more "natural" implementation
index fe3a55e..766ca11 100644 (file)
@@ -629,6 +629,7 @@ alter = go
 -- In short : @'lookup' k <$> 'alterF' f k m = f ('lookup' k m)@.
 --
 -- Example:
+--
 -- @
 -- interactiveAlter :: Int -> Map Int String -> IO (Map Int String)
 -- interactiveAlter k m = alterF f k m where
@@ -651,13 +652,12 @@ alter = go
 --
 -- Note on rewrite rules:
 --
--- This module includes GHC rewrite rules to optimize 'alterF' for
--- the 'Const' and 'Identity' functors. In general, these rules
--- improve performance. The sole exception is that when using
--- 'Identity', deleting a key that is already absent takes longer
--- than it would without the rules. If you expect this to occur
--- a very large fraction of the time, you might consider using a
--- private copy of the 'Identity' type.
+-- This module includes GHC rewrite rules to optimize 'alterF' for the 'Const',
+-- 'Identity', and @(,) b@ functors. In general, these rules improve
+-- performance. The main exception is that when using 'Identity', deleting a
+-- key that is already absent takes longer than it would without the rules. If
+-- you expect this to occur a very large fraction of the time, you might
+-- consider using a private copy of the 'Identity' type.
 --
 -- Note: 'alterF' is a flipped version of the 'at' combinator from
 -- 'Control.Lens.At'.
@@ -674,7 +674,13 @@ alterF f k m = atKeyImpl Strict k f m
 -- `Control.Applicative.Const` and just doing a lookup.
 {-# RULES
 "alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m
+"alterF/Pair" forall k (f :: Maybe a -> (b, Maybe a)) . alterF f k = atKeyPair k f
  #-}
+
+atKeyPair :: Ord k => k -> (Maybe a -> (b, Maybe a)) -> Map k a -> (b, Map k a)
+atKeyPair k f t = atKeyWithLookup Strict k f t
+{-# INLINABLE atKeyPair #-}
+
 #if MIN_VERSION_base(4,8,0)
 -- base 4.8 and above include Data.Functor.Identity, so we can
 -- save a pretty decent amount of time by handling it specially.
index f0ba0b4..3939bb8 100644 (file)
@@ -24,6 +24,18 @@ main = do
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
+        , bench "updateLookupWithKey absent" $ whnf (upd' Just evens) m_odd
+        , bench "updateLookupWithKey present" $ whnf (upd' Just evens) m_even
+        , bench "updateLookupWithKey delete" $ whnf (upd' (const Nothing) evens) m
+        , bench "alterF updateLookupWithKey absent" $ whnf (updlAlterF Just evens) m_odd
+        , bench "alterF updateLookupWithKey present" $ whnf (updlAlterF Just evens) m_even
+        , bench "alterF updateLookupWithKey delete" $ whnf (updlAlterF (const Nothing) evens) m
+        , bench "insertLookupWithKey absent" $ whnf (insLookupWithKey elems_even) m_odd
+        , bench "insertLookupWithKey present" $ whnf (insLookupWithKey elems_even) m_even
+        , bench "insertLookupWithKeyAlterF absent" $ whnf (insLookupWithKeyAlterF elems_even) m_odd
+        , bench "insertLookupWithKeyAlterF present" $ whnf (insLookupWithKeyAlterF elems_even) m_even
+        , bench "insertLookupWithKey' absent" $ whnf (insLookupWithKey' elems_even) m_odd
+        , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
         , 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
@@ -60,10 +72,6 @@ main = do
         , bench "insertWithKey present" $ whnf (insWithKey elems_even) m_even
         , bench "insertWithKey' absent" $ whnf (insWithKey' elems_even) m_odd
         , bench "insertWithKey' present" $ whnf (insWithKey' elems_even) m_even
-        , bench "insertLookupWithKey absent" $ whnf (insLookupWithKey elems_even) m_odd
-        , bench "insertLookupWithKey present" $ whnf (insLookupWithKey elems_even) m_even
-        , bench "insertLookupWithKey' absent" $ whnf (insLookupWithKey' elems_even) m_odd
-        , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
         , bench "map" $ whnf (M.map (+ 1)) m
         , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
         , bench "foldlWithKey" $ whnf (ins elems) m
@@ -72,9 +80,6 @@ main = do
         , bench "update absent" $ whnf (upd Just evens) m_odd
         , bench "update present" $ whnf (upd Just evens) m_even
         , bench "update delete" $ whnf (upd (const Nothing) evens) m
-        , bench "updateLookupWithKey absent" $ whnf (upd' Just evens) m_odd
-        , bench "updateLookupWithKey present" $ whnf (upd' Just evens) m_even
-        , bench "updateLookupWithKey delete" $ whnf (upd' (const Nothing) evens) m
         , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
         , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
         , bench "lookupIndex" $ whnf (lookupIndex keys) m
@@ -156,6 +161,18 @@ insLookupWithKey xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b)
     f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey add3 k v m
                         in PS (fromMaybe 0 n' + n) m'
 
+insLookupWithKeyAlterF :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int)
+insLookupWithKeyAlterF xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b)
+  where
+    f (PS n m) (k, v) = let !(n', m') = insertLookupWithKeyAlterF add3 k v m
+                        in PS (fromMaybe 0 n' + n) m'
+
+insertLookupWithKeyAlterF :: Ord k => (k -> a -> a -> a) -> k -> a -> M.Map k a
+                    -> (Maybe a, M.Map k a)
+insertLookupWithKeyAlterF f k y m = M.alterF go k m where
+  go Nothing = (Nothing, Just y)
+  go old@(Just x) = (old, Just (f k y x))
+
 insLookupWithKey' :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int)
 insLookupWithKey' xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b)
   where
@@ -177,6 +194,14 @@ upd f xs m = foldl' (\m k -> M.update f k m) m xs
 upd' :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
 upd' f xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> f a) k m) m xs
 
+updlAlterF :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
+updlAlterF f xs m = foldl' (\m k -> snd $ updateLookupWithKeyAlterF (\_ a -> f a) k m) m xs
+
+updateLookupWithKeyAlterF :: Ord k => (k -> a -> Maybe a) -> k -> M.Map k a -> (Maybe a,M.Map k a)
+updateLookupWithKeyAlterF f k m = M.alterF go k m where
+  go Nothing = (Nothing, Nothing)
+  go o@(Just old) = (o, f k old)
+
 alt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
 alt f xs m = foldl' (\m k -> M.alter f k m) m xs