Replace atomicModifyMutVar#

Reviewers: simonmar, hvr, bgamari, erikd, fryguybob, rrnewton

Reviewed By: simonmar

Subscribers: fryguybob, rwbarton, thomie, carter

GHC Trac Issues: #15364

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

7 files changed, 168 insertions, 33 deletions

diff --git a/libraries/base/Data/IORef.hs b/libraries/base/Data/IORef.hs
index c6275f5433..6f07a097f7 100644
--- a/libraries/base/Data/IORef.hs
+++ b/libraries/base/Data/IORef.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples #-}
+{-# LANGUAGE BangPatterns #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -36,8 +37,7 @@ module Data.IORef
 
 import GHC.Base
 import GHC.STRef
-import GHC.IORef hiding (atomicModifyIORef)
-import qualified GHC.IORef
+import GHC.IORef
 import GHC.Weak
 
 -- |Make a 'Weak' pointer to an 'IORef', using the second argument as a finalizer
@@ -91,18 +91,9 @@ modifyIORef' ref f = do
 -- Use 'atomicModifyIORef'' or 'atomicWriteIORef' to avoid this problem.
 --
 atomicModifyIORef :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef = GHC.IORef.atomicModifyIORef
-
--- | Strict version of 'atomicModifyIORef'.  This forces both the value stored
--- in the 'IORef' as well as the value returned.
---
--- @since 4.6.0.0
-atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef' ref f = do
-    b <- atomicModifyIORef ref $ \a ->
-            case f a of
-                v@(a',_) -> a' `seq` v
-    b `seq` return b
+atomicModifyIORef ref f = do
+  (_old, ~(_new, res)) <- atomicModifyIORef2 ref f
+  pure res
 
 -- | Variant of 'writeIORef' with the \"barrier to reordering\" property that
 -- 'atomicModifyIORef' has.
@@ -110,8 +101,8 @@ atomicModifyIORef' ref f = do
 -- @since 4.6.0.0
 atomicWriteIORef :: IORef a -> a -> IO ()
 atomicWriteIORef ref a = do
-    x <- atomicModifyIORef ref (\_ -> (a, ()))
-    x `seq` return ()
+  _ <- atomicSwapIORef ref a
+  pure ()
 
 {- $memmodel
 
diff --git a/libraries/base/GHC/Conc/Windows.hs b/libraries/base/GHC/Conc/Windows.hs
index 6c4136428a..ed5e0452a0 100644
--- a/libraries/base/GHC/Conc/Windows.hs
+++ b/libraries/base/GHC/Conc/Windows.hs
@@ -131,7 +131,7 @@ waitForDelayEvent :: Int -> IO ()
 waitForDelayEvent usecs = do
   m <- newEmptyMVar
   target <- calculateTarget usecs
-  atomicModifyIORef pendingDelays (\xs -> (Delay target m : xs, ()))
+  _ <- atomicModifyIORef'_ pendingDelays (\xs -> Delay target m : xs)
   prodServiceThread
   takeMVar m
 
@@ -140,7 +140,7 @@ waitForDelayEventSTM :: Int -> IO (TVar Bool)
 waitForDelayEventSTM usecs = do
    t <- atomically $ newTVar False
    target <- calculateTarget usecs
-   atomicModifyIORef pendingDelays (\xs -> (DelaySTM target t : xs, ()))
+   _ <- atomicModifyIORef'_ pendingDelays (\xs -> DelaySTM target t : xs)
    prodServiceThread
    return t
 
@@ -219,10 +219,10 @@ foreign import ccall unsafe "getOrSetGHCConcWindowsProddingStore"
 
 prodServiceThread :: IO ()
 prodServiceThread = do
-  -- NB. use atomicModifyIORef here, otherwise there are race
+  -- NB. use atomicSwapIORef here, otherwise there are race
   -- conditions in which prodding is left at True but the server is
   -- blocked in select().
-  was_set <- atomicModifyIORef prodding $ \b -> (True,b)
+  was_set <- atomicSwapIORef prodding True
   when (not was_set) wakeupIOManager
 
 -- ----------------------------------------------------------------------------
@@ -239,7 +239,7 @@ service_loop :: HANDLE          -- read end of pipe
 
 service_loop wakeup old_delays = do
   -- pick up new delay requests
-  new_delays <- atomicModifyIORef pendingDelays (\a -> ([],a))
+  new_delays <- atomicSwapIORef pendingDelays []
   let  delays = foldr insertDelay old_delays new_delays
 
   now <- getMonotonicUSec
@@ -262,8 +262,7 @@ service_loop wakeup old_delays = do
 
 service_cont :: HANDLE -> [DelayReq] -> IO ()
 service_cont wakeup delays = do
-  r <- atomicModifyIORef prodding (\_ -> (False,False))
-  r `seq` return () -- avoid space leak
+  _ <- atomicSwapIORef prodding False
   service_loop wakeup delays
 
 -- must agree with rts/win32/ThrIOManager.c
diff --git a/libraries/base/GHC/Event/Control.hs b/libraries/base/GHC/Event/Control.hs
index 5b4a81b38b..779d60d5d7 100644
--- a/libraries/base/GHC/Event/Control.hs
+++ b/libraries/base/GHC/Event/Control.hs
@@ -126,7 +126,7 @@ newControl shouldRegister = allocaArray 2 $ \fds -> do
 -- file after it has been closed.
 closeControl :: Control -> IO ()
 closeControl w = do
-  atomicModifyIORef (controlIsDead w) (\_ -> (True, ()))
+  _ <- atomicSwapIORef (controlIsDead w) True
   _ <- c_close . fromIntegral . controlReadFd $ w
   _ <- c_close . fromIntegral . controlWriteFd $ w
   when (didRegisterWakeupFd w) $ c_setIOManagerWakeupFd (-1)
diff --git a/libraries/base/GHC/Exts.hs b/libraries/base/GHC/Exts.hs
index 11d329ac19..3f5b630ab9 100755
--- a/libraries/base/GHC/Exts.hs
+++ b/libraries/base/GHC/Exts.hs
@@ -34,6 +34,9 @@ module GHC.Exts
         uncheckedIShiftL64#, uncheckedIShiftRA64#,
         isTrue#,
 
+        -- * Compat wrapper
+        atomicModifyMutVar#,
+
         -- * Fusion
         build, augment,
 
@@ -219,3 +222,27 @@ instance IsList CallStack where
   type (Item CallStack) = (String, SrcLoc)
   fromList = fromCallSiteList
   toList   = getCallStack
+
+-- | An implementation of the old @atomicModifyMutVar#@ primop in
+-- terms of the new 'atomicModifyMutVar2#' primop, for backwards
+-- compatibility. The type of this function is a bit bogus. It's
+-- best to think of it as having type
+--
+-- @
+-- atomicModifyMutVar#
+--   :: MutVar# s a
+--   -> (a -> (a, b))
+--   -> State# s
+--   -> (# State# s, b #)
+-- @
+--
+-- but there may be code that uses this with other two-field record
+-- types.
+atomicModifyMutVar#
+  :: MutVar# s a
+  -> (a -> b)
+  -> State# s
+  -> (# State# s, c #)
+atomicModifyMutVar# mv f s =
+  case unsafeCoerce# (atomicModifyMutVar2# mv f s) of
+    (# s', _, ~(_, res) #) -> (# s', res #)
diff --git a/libraries/base/GHC/ForeignPtr.hs b/libraries/base/GHC/ForeignPtr.hs
index 043de1f94b..d80412aac1 100644
--- a/libraries/base/GHC/ForeignPtr.hs
+++ b/libraries/base/GHC/ForeignPtr.hs
@@ -321,7 +321,7 @@ addForeignPtrConcFinalizer_ _ _ =
 
 insertHaskellFinalizer :: IORef Finalizers -> IO () -> IO Bool
 insertHaskellFinalizer r f = do
-  !wasEmpty <- atomicModifyIORef r $ \finalizers -> case finalizers of
+  !wasEmpty <- atomicModifyIORefP r $ \finalizers -> case finalizers of
       NoFinalizers -> (HaskellFinalizers [f], True)
       HaskellFinalizers fs -> (HaskellFinalizers (f:fs), False)
       _ -> noMixingError
@@ -352,8 +352,8 @@ ensureCFinalizerWeak ref@(IORef (STRef r#)) value = do
       NoFinalizers -> IO $ \s ->
           case mkWeakNoFinalizer# r# (unsafeCoerce# value) s of { (# s1, w #) ->
              -- See Note [MallocPtr finalizers] (#10904)
-          case atomicModifyMutVar# r# (update w) s1 of
-              { (# s2, (weak, needKill ) #) ->
+          case atomicModifyMutVar2# r# (update w) s1 of
+              { (# s2, _, (_, (weak, needKill )) #) ->
           if needKill
             then case finalizeWeak# w s2 of { (# s3, _, _ #) ->
               (# s3, weak #) }
@@ -370,7 +370,8 @@ noMixingError = errorWithoutStackTrace $
 
 foreignPtrFinalizer :: IORef Finalizers -> IO ()
 foreignPtrFinalizer r = do
-  fs <- atomicModifyIORef r $ \fs -> (NoFinalizers, fs) -- atomic, see #7170
+  fs <- atomicSwapIORef r NoFinalizers
+             -- atomic, see #7170
   case fs of
     NoFinalizers -> return ()
     CFinalizers w -> IO $ \s -> case finalizeWeak# w s of
diff --git a/libraries/base/GHC/IORef.hs b/libraries/base/GHC/IORef.hs
index 7377690f0f..d04ae728fd 100644
--- a/libraries/base/GHC/IORef.hs
+++ b/libraries/base/GHC/IORef.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE Unsafe #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash #-}
-{-# OPTIONS_GHC -funbox-strict-fields #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE BangPatterns #-}
 {-# OPTIONS_HADDOCK hide #-}
 
 -----------------------------------------------------------------------------
@@ -19,7 +20,9 @@
 
 module GHC.IORef (
         IORef(..),
-        newIORef, readIORef, writeIORef, atomicModifyIORef
+        newIORef, readIORef, writeIORef, atomicModifyIORef2Lazy,
+        atomicModifyIORef2, atomicModifyIORefLazy_, atomicModifyIORef'_,
+        atomicModifyIORefP, atomicSwapIORef, atomicModifyIORef'
     ) where
 
 import GHC.Base
@@ -48,6 +51,120 @@ readIORef  (IORef var) = stToIO (readSTRef var)
 writeIORef  :: IORef a -> a -> IO ()
 writeIORef (IORef var) v = stToIO (writeSTRef var v)
 
-atomicModifyIORef :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef (IORef (STRef r#)) f = IO $ \s -> atomicModifyMutVar# r# f s
+-- Atomically apply a function to the contents of an 'IORef',
+-- installing its first component in the 'IORef' and returning
+-- the old contents and the result of applying the function.
+-- The result of the function application (the pair) is not forced.
+-- As a result, this can lead to memory leaks. It is generally better
+-- to use 'atomicModifyIORef2'.
+atomicModifyIORef2Lazy :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
+atomicModifyIORef2Lazy (IORef (STRef r#)) f =
+  IO (\s -> case atomicModifyMutVar2# r# f s of
+              (# s', old, res #) -> (# s', (old, res) #))
 
+-- Atomically apply a function to the contents of an 'IORef',
+-- installing its first component in the 'IORef' and returning
+-- the old contents and the result of applying the function.
+-- The result of the function application (the pair) is forced,
+-- but neither of its components is.
+atomicModifyIORef2 :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
+atomicModifyIORef2 ref f = do
+  r@(_old, (_new, _res)) <- atomicModifyIORef2Lazy ref f
+  return r
+
+-- | A version of 'Data.IORef.atomicModifyIORef' that forces
+-- the (pair) result of the function.
+atomicModifyIORefP :: IORef a -> (a -> (a,b)) -> IO b
+atomicModifyIORefP ref f = do
+  (_old, (_,r)) <- atomicModifyIORef2 ref f
+  pure r
+
+-- | Atomically apply a function to the contents of an
+-- 'IORef' and return the old and new values. The result
+-- of the function is not forced. As this can lead to a
+-- memory leak, it is usually better to use `atomicModifyIORef'_`.
+atomicModifyIORefLazy_ :: IORef a -> (a -> a) -> IO (a, a)
+atomicModifyIORefLazy_ (IORef (STRef ref)) f = IO $ \s ->
+  case atomicModifyMutVar_# ref f s of
+    (# s', old, new #) -> (# s', (old, new) #)
+
+-- | Atomically apply a function to the contents of an
+-- 'IORef' and return the old and new values. The result
+-- of the function is forced.
+atomicModifyIORef'_ :: IORef a -> (a -> a) -> IO (a, a)
+atomicModifyIORef'_ ref f = do
+  (old, !new) <- atomicModifyIORefLazy_ ref f
+  return (old, new)
+
+-- | Atomically replace the contents of an 'IORef', returning
+-- the old contents.
+atomicSwapIORef :: IORef a -> a -> IO a
+-- Bad implementation! This will be a primop shortly.
+atomicSwapIORef (IORef (STRef ref)) new = IO $ \s ->
+  case atomicModifyMutVar2# ref (\_old -> Box new) s of
+    (# s', old, Box _new #) -> (# s', old #)
+
+data Box a = Box a
+
+-- | Strict version of 'Data.IORef.atomicModifyIORef'. This forces both
+-- the value stored in the 'IORef' and the value returned. The new value
+-- is installed in the 'IORef' before the returned value is forced.
+-- So
+--
+-- @atomicModifyIORef' ref (\x -> (x+1, undefined))@
+--
+-- will increment the 'IORef' and then throw an exception in the calling
+-- thread.
+--
+-- @since 4.6.0.0
+atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
+-- See Note [atomicModifyIORef' definition]
+atomicModifyIORef' ref f = do
+  (_old, (_new, !res)) <- atomicModifyIORef2 ref $
+    \old -> case f old of
+       r@(!_new, _res) -> r
+  pure res
+
+-- Note [atomicModifyIORef' definition]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- atomicModifyIORef' was historically defined
+--
+--    atomicModifyIORef' ref f = do
+--        b <- atomicModifyIORef ref $ \a ->
+--                case f a of
+--                    v@(a',_) -> a' `seq` v
+--        b `seq` return b
+--
+-- The most obvious definition, now that we have atomicModifyMutVar2#,
+-- would be
+--
+--    atomicModifyIORef' ref f = do
+--      (_old, (!_new, !res)) <- atomicModifyIORef2 ref f
+--      pure res
+--
+-- Why do we force the new value on the "inside" instead of afterwards?
+-- I initially thought the latter would be okay, but then I realized
+-- that if we write
+--
+--   atomicModifyIORef' ref $ \x -> (x + 5, x - 5)
+--
+-- then we'll end up building a pair of thunks to calculate x + 5
+-- and x - 5. That's no good! With the more complicated definition,
+-- we avoid this problem; the result pair is strict in the new IORef
+-- contents. Of course, if the function passed to atomicModifyIORef'
+-- doesn't inline, we'll build a closure for it. But that was already
+-- true for the historical definition of atomicModifyIORef' (in terms
+-- of atomicModifyIORef), so we shouldn't lose anything. Note that
+-- in keeping with the historical behavior, we *don't* propagate the
+-- strict demand on the result inwards. In particular,
+--
+--   atomicModifyIORef' ref (\x -> (x + 1, undefined))
+--
+-- will increment the IORef and throw an exception; it will not
+-- install an undefined value in the IORef.
+--
+-- A clearer version, in my opinion (but one quite incompatible with
+-- the traditional one) would only force the new IORef value and not
+-- the result. This version would have been relatively inefficient
+-- to implement using atomicModifyMutVar#, but is just fine now.
diff --git a/libraries/base/System/IO.hs b/libraries/base/System/IO.hs
index e02c30d63c..4a9b8837ab 100644
--- a/libraries/base/System/IO.hs
+++ b/libraries/base/System/IO.hs
@@ -583,7 +583,7 @@ tempCounter = unsafePerformIO $ newIORef 0
 rand_string :: IO String
 rand_string = do
   r1 <- c_getpid
-  r2 <- atomicModifyIORef tempCounter (\n -> (n+1, n))
+  (r2, _) <- atomicModifyIORef'_ tempCounter (+1)
   return $ show r1 ++ "-" ++ show r2
 
 data OpenNewFileResult