Improve GHC.Event.IntTable performance

Speed up GHC.Event.IntTable.lookup by removing the IO context from the
go helper function. This generates a little bit better code as we can
avoid repeating the stack check.

Remove unused parameter from GHC.Event.IntTable.updateWith.go and
directly return a bool instead of a maybe and then checking that whether
it is a Nothing.

Test Plan: validate

Reviewers: austin, hvr, bgamari

Reviewed By: bgamari

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1742

GHC Trac Issues: #8793

1 files changed, 15 insertions, 16 deletions

diff --git a/libraries/base/GHC/Event/IntTable.hs b/libraries/base/GHC/Event/IntTable.hs
index ea487d5391..7ae2e1a8dc 100644
--- a/libraries/base/GHC/Event/IntTable.hs
+++ b/libraries/base/GHC/Event/IntTable.hs
@@ -15,10 +15,10 @@ module GHC.Event.IntTable
 
 import Data.Bits ((.&.), shiftL, shiftR)
 import Data.IORef (IORef, newIORef, readIORef, writeIORef)
-import Data.Maybe (Maybe(..), isJust, isNothing)
+import Data.Maybe (Maybe(..), isJust)
 import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtr, withForeignPtr)
 import Foreign.Storable (peek, poke)
-import GHC.Base (Monad(..), (=<<), ($), const, liftM, otherwise, when)
+import GHC.Base (Monad(..), (=<<), ($), ($!), const, liftM, otherwise, when)
 import GHC.Classes (Eq(..), Ord(..))
 import GHC.Event.Arr (Arr)
 import GHC.Num (Num(..))
@@ -47,11 +47,12 @@ data Bucket a = Empty
 lookup :: Int -> IntTable a -> IO (Maybe a)
 lookup k (IntTable ref) = do
   let go Bucket{..}
-        | bucketKey == k = return (Just bucketValue)
+        | bucketKey == k = Just bucketValue
         | otherwise      = go bucketNext
-      go _ = return Nothing
+      go _ = Nothing
   it@IT{..} <- readIORef ref
-  go =<< Arr.read tabArr (indexOf k it)
+  bkt <- Arr.read tabArr (indexOf k it)
+  return $! go bkt
 
 new :: Int -> IO (IntTable a)
 new capacity = IntTable `liftM` (newIORef =<< new_ capacity)
@@ -125,20 +126,18 @@ updateWith :: (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
 updateWith f k (IntTable ref) = do
   it@IT{..} <- readIORef ref
   let idx = indexOf k it
-      go changed bkt@Bucket{..}
-        | bucketKey == k =
-            let fbv = f bucketValue
-                !nb = case fbv of
-                        Just val -> bkt { bucketValue = val }
-                        Nothing  -> bucketNext
-            in (fbv, Just bucketValue, nb)
-        | otherwise = case go changed bucketNext of
+      go bkt@Bucket{..}
+        | bucketKey == k = case f bucketValue of
+            Just val -> let !nb = bkt { bucketValue = val }
+                        in (False, Just bucketValue, nb)
+            Nothing  -> (True, Just bucketValue, bucketNext)
+        | otherwise = case go bucketNext of
                         (fbv, ov, nb) -> (fbv, ov, bkt { bucketNext = nb })
-      go _ e = (Nothing, Nothing, e)
-  (fbv, oldVal, newBucket) <- go False `liftM` Arr.read tabArr idx
+      go e = (False, Nothing, e)
+  (del, oldVal, newBucket) <- go `liftM` Arr.read tabArr idx
   when (isJust oldVal) $ do
     Arr.write tabArr idx newBucket
-    when (isNothing fbv) $
+    when del $
       withForeignPtr tabSize $ \ptr -> do
         size <- peek ptr
         poke ptr (size - 1)