From 1f61c772f9d19c2f40e37166224003c5fb65c4e9 Mon Sep 17 00:00:00 2001
From: Matt Renaud <matt@m-renaud.com>
Date: Wed, 31 Jan 2018 16:59:38 -0800
Subject: [PATCH] Optimize IntMap.alter using unboxed sums.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We use unboxed sums (available in GHC >= 8.2) to track whether or not the map
required modification.

Benchmark results:
------------------

Before:
-------

benchmarking alter
time                 881.1 μs   (862.0 μs .. 899.2 μs)
                     0.996 R²   (0.994 R² .. 0.998 R²)
mean                 854.9 μs   (841.8 μs .. 872.5 μs)
std dev              49.55 μs   (38.01 μs .. 69.97 μs)
variance introduced by outliers: 48% (moderately inflated)

After:
------

benchmarking alter
time                 513.1 μs   (506.5 μs .. 519.6 μs)
                     0.998 R²   (0.997 R² .. 0.999 R²)
mean                 517.8 μs   (511.7 μs .. 528.5 μs)
std dev              27.05 μs   (17.74 μs .. 47.52 μs)
variance introduced by outliers: 45% (moderately inflated)
---
 Data/IntMap/Internal.hs | 50 +++++++++++++++++++++++++++++++++++++++++
 Data/IntMap/Strict.hs   | 49 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/Data/IntMap/Internal.hs b/Data/IntMap/Internal.hs
index 9b30673cc..699ab57e7 100644
--- a/Data/IntMap/Internal.hs
+++ b/Data/IntMap/Internal.hs
@@ -11,6 +11,10 @@
 #if __GLASGOW_HASKELL__ >= 708
 {-# LANGUAGE TypeFamilies #-}
 #endif
+#if __GLASGOW_HASKELL__ >= 802
+{-# LANGUAGE UnboxedSums #-}
+{-# LANGUAGE UnboxedTuples #-}
+#endif
 
 {-# OPTIONS_HADDOCK not-home #-}
 
@@ -311,6 +315,9 @@ import Prelude hiding (lookup, map, filter, foldr, foldl, null)
 import Data.IntSet.Internal (Key)
 import qualified Data.IntSet.Internal as IntSet
 import Utils.Containers.Internal.BitUtil
+#if __GLASGOW_HASKELL__ >= 802
+import Utils.Containers.Internal.PtrEquality (ptrEq)
+#endif
 import Utils.Containers.Internal.StrictFold
 import Utils.Containers.Internal.StrictPair
 
@@ -937,6 +944,48 @@ updateLookupWithKey _ _ Nil = (Nothing,Nil)
 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
 -- 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'
+{-# INLINE alter #-}
+
+
+-- Internal implementation which keeps track of whether or not the intmap was
+-- modified using an unboxed sum (Maybe).
+--
+-- 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# 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 #)
+  | zero k m = case alter# f k l of
+      (# (# #) | #) -> (# (# #) | #)
+      (# | l' #) -> (# | binCheckLeft p m l' r #)
+
+  | otherwise = case alter# f k r of
+      (# (# #) | #) -> (# (# #) | #)
+      (# | r' #) -> (# | 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 -> (# (# #) | #)
+{-# INLINABLE alter# #-}
+
+#else
 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
 alter f !k t@(Bin p m l r)
   | nomatch k p m = case f Nothing of
@@ -954,6 +1003,7 @@ alter f k t@(Tip ky y)
 alter f k Nil     = case f Nothing of
                       Just x -> Tip k x
                       Nothing -> Nil
+#endif
 
 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
 -- @k@, or absence thereof.  'alterF' can be used to inspect, insert, delete,
diff --git a/Data/IntMap/Strict.hs b/Data/IntMap/Strict.hs
index dc638f023..30c45b8a7 100644
--- a/Data/IntMap/Strict.hs
+++ b/Data/IntMap/Strict.hs
@@ -3,6 +3,11 @@
 #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
 {-# LANGUAGE Trustworthy #-}
 #endif
+#if __GLASGOW_HASKELL__ >= 802
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedSums #-}
+{-# LANGUAGE UnboxedTuples #-}
+#endif
 
 #include "containers.h"
 
@@ -303,6 +308,9 @@ import Data.IntMap.Internal
 import Data.IntMap.Internal.DeprecatedDebug (showTree, showTreeWith)
 import qualified Data.IntSet.Internal as IntSet
 import Utils.Containers.Internal.BitUtil
+#if __GLASGOW_HASKELL__ >= 802
+import Utils.Containers.Internal.PtrEquality (ptrEq)
+#endif
 import Utils.Containers.Internal.StrictFold
 import Utils.Containers.Internal.StrictPair
 #if !MIN_VERSION_base(4,8,0)
@@ -559,6 +567,46 @@ updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0
 -- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
 -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'.
 -- 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'
+{-# INLINE alter #-}
+
+
+-- Internal implementation which keeps track of whether or not the intmap was
+-- modified using an unboxed sum (Maybe).
+--
+-- 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# 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 #)
+  | zero k m = case alter# f k l of
+      (# (# #) | #) -> (# (# #) | #)
+      (# | l' #) -> (# | binCheckLeft p m l' r #)
+
+  | otherwise = case alter# f k r of
+      (# (# #) | #) -> (# (# #) | #)
+      (# | r' #) -> (# | 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 -> (# (# #) | #)
+
+#else
 alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
 alter f !k t =
   case t of
@@ -578,6 +626,7 @@ alter f !k t =
     Nil               -> case f Nothing of
                            Just !x -> Tip k x
                            Nothing -> Nil
+#endif
 
 -- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at
 -- @k@, or absence thereof.  'alterF' can be used to inspect, insert, delete,