Strictify, simplify

Unboxed sums are lazy in their lifted fields, so we need to
be explicitly strict. Using pattern synonyms cuts down
tremendously on the noise.

Author: treeowl

Data/IntMap/Strict.hs

@@ -7,6 +7,9 @@
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE UnboxedSums #-}
 {-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeInType #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 #endif
 
 #include "containers.h"
@@ -221,7 +224,7 @@ module Data.IntMap.Strict (
     , showTreeWith
     ) where
 
-import Prelude hiding (lookup,map,filter,foldr,foldl,null)
+import Prelude hiding (lookup,map,filter,foldr,foldl,null, ($!))
 
 import Data.Bits
 import qualified Data.IntMap.Internal as L
@@ -310,6 +313,8 @@ import qualified Data.IntSet.Internal as IntSet
 import Utils.Containers.Internal.BitUtil
 #if __GLASGOW_HASKELL__ >= 802
 import Utils.Containers.Internal.PtrEquality (ptrEq)
+import Utils.Containers.Internal.UnboxedMaybe (Maybe#, pattern Nothing#, pattern Just#, toMaybe)
+import GHC.Exts (TYPE)
 #endif
 import Utils.Containers.Internal.StrictFold
 import Utils.Containers.Internal.StrictPair
@@ -336,6 +341,19 @@ import Control.Applicative (Applicative (..), liftA2)
 -- > map (\ v -> undefined) m  ==  undefined      -- m is not empty
 -- > mapKeys (\ k -> undefined) m  ==  undefined  -- m is not empty
 
+#if __GLASGOW_HASKELL__ >= 802
+-- Annoyingly, Prelude.$! is not (yet) polykinded. So we just define
+-- our own.
+infixr 0 $!
+($!) :: forall r a (b :: TYPE r).
+        (a -> b) -> a -> b
+f $! (!x) = f x
+#endif
+#else
+($!) :: (a -> b) -> a -> b
+f $! (!x) = f x
+#endif
+
 {--------------------------------------------------------------------
   Query
 --------------------------------------------------------------------}
@@ -569,9 +587,11 @@ updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0
 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
 #if __GLASGOW_HASKELL__ >= 802
 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
-alter f !k t = case alter# f k t of
-  (# (# #) | #) -> t
-  (# | t' #) -> t'
+alter f !k t = case alter# (\m -> case f (toMaybe m) of
+                              Nothing -> Nothing#
+                              Just x -> Just# x) k t of
+  Nothing# -> t
+  Just# t' -> t'
 {-# INLINE alter #-}
 
 
@@ -580,31 +600,31 @@ alter f !k t = case alter# f k t of
 --
 -- If no modifications are made to the map (# (# #) | #) is returned, otherwise
 -- (# | newMap #) is returned.
-alter# :: (Maybe a -> Maybe a) -> Key -> IntMap a -> (# (# #) | IntMap a #)
+alter# :: (Maybe# a -> Maybe# a) -> Key -> IntMap a -> Maybe# (IntMap a)
 alter# f !k t@(Bin p m l r)
-  | nomatch k p m = case f Nothing of
-                      Nothing -> (# (# #) | #)
-                      Just !x -> (# | link k (Tip k x) p t #)
+  | nomatch k p m = case f Nothing# of
+                      Nothing# -> Nothing#
+                      Just# !x -> Just# $! link k (Tip k x) p t
   | zero k m = case alter# f k l of
-      (# (# #) | #) -> (# (# #) | #)
-      (# | l' #) -> (# | binCheckLeft p m l' r #)
+      Nothing# -> Nothing#
+      Just# l' -> Just# $! binCheckLeft p m l' r
 
   | otherwise = case alter# f k r of
-      (# (# #) | #) -> (# (# #) | #)
-      (# | r' #) -> (# | binCheckRight p m l r' #)
+      Nothing# -> Nothing#
+      Just# r' -> Just# $! binCheckRight p m l r'
 
 alter# f k t@(Tip ky y)
-  | k==ky         = case f (Just y) of
-                      Just x -> if x `ptrEq` y
-                                then (# (# #) | #)
-                                else (# | Tip ky x #)
-                      Nothing -> (# | Nil #)
-  | otherwise     = case f Nothing of
-                      Just !x -> (# | link k (Tip k x) ky t #)
-                      Nothing -> (# (# #) | #)
-alter# f k Nil     = case f Nothing of
-                      Just !x -> (# | Tip k x #)
-                      Nothing -> (# (# #) | #)
+  | k==ky         = case f (Just# y) of
+                      Just# !x -> if x `ptrEq` y
+                                  then Nothing#
+                                  else Just# (Tip ky x)
+                      Nothing# -> Just# Nil
+  | otherwise     = case f Nothing# of
+                      Just# !x -> Just# $! link k (Tip k x) ky t
+                      Nothing# -> Nothing#
+alter# f k Nil     = case f Nothing# of
+                      Just# !x -> Just# (Tip k x)
+                      Nothing# -> Nothing#
 
 #else
 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a

Utils/Containers/Internal/UnboxedMaybe.hs

@@ -0,0 +1,19 @@
+{-# language MagicHash, UnboxedSums, UnboxedTuples, PatternSynonyms #-}
+
+module Utils.Containers.Internal.UnboxedMaybe (Maybe#, pattern Nothing#, pattern Just#,
+  toMaybe) where
+
+type Maybe# a = (# (# #) | a #)
+
+pattern Nothing# :: Maybe# a
+pattern Nothing# = (# (# #) | #)
+
+pattern Just# :: a -> Maybe# a
+pattern Just# a = (# | a #)
+
+{-# COMPLETE Nothing#, Just# #-}
+
+toMaybe :: Maybe# a -> Maybe a
+toMaybe (Just# a) = Just a
+toMaybe _ = Nothing
+{-# INLINE toMaybe #-}

benchmarks/IntMap.hs

@@ -34,7 +34,13 @@ main = do
         , bench "delete" $ whnf (del keys) m
         , bench "update" $ whnf (upd keys) m
         , bench "updateLookupWithKey" $ whnf (upd' keys) m
-        , bench "alter"  $ whnf (alt keys) m
+        , bench "alter"  $ whnf (alt id keys) m
+
+        , bench "alter absent"  $ whnf (alt id evens) m_odd
+        , bench "alter insert"  $ whnf (alt (const (Just 1)) evens) m_odd
+        , bench "alter update"  $ whnf (alt id evens) m_even
+        , bench "alter delete"  $ whnf (alt (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
@@ -49,6 +55,13 @@ main = do
     values = [1..2^12]
     sum k v1 v2 = k + v1 + v2
     consPair k v xs = (k, v) : xs
+    m_even = M.fromAscList elems_even :: M.IntMap Int
+    m_odd = M.fromAscList elems_odd :: M.IntMap Int
+    elems_even = zip evens evens
+    elems_odd = zip odds odds
+    evens = [2,4..bound]
+    odds = [1,3..bound]
+    bound = 2^12
 
 add3 :: Int -> Int -> Int -> Int
 add3 x y z = x + y + z
@@ -89,8 +102,8 @@ upd xs m = foldl' (\m k -> M.update Just k m) m xs
 upd' :: [Int] -> M.IntMap Int -> M.IntMap Int
 upd' xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> Just a) k m) m xs
 
-alt :: [Int] -> M.IntMap Int -> M.IntMap Int
-alt xs m = foldl' (\m k -> M.alter id k m) m xs
+alt :: (Maybe Int -> Maybe Int) -> [Int] -> M.IntMap Int -> M.IntMap Int
+alt f xs m = foldl' (\m k -> M.alter f k m) m xs
 
 maybeDel :: Int -> Maybe Int
 maybeDel n | n `mod` 3 == 0 = Nothing

containers.cabal

@@ -88,6 +88,9 @@ Library
         Utils.Containers.Internal.Coercions
         Data.Map.Internal.DeprecatedShowTree
         Data.IntMap.Internal.DeprecatedDebug
+    if impl (ghc >= 7.2)
+      other-modules:
+        Utils.Containers.Internal.UnboxedMaybe
 
     include-dirs: include