{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE UnboxedTuples #-}
-- | Atomic 'IORef' and array modification operations for more general result
-- types.
module Data.IORef.AtomicModify
( atomicModifyIORef3General
, atomicModifyArray3General
, atomicModifySmallArray3General
) where
import GHC.IORef (IORef (..), newIORef, readIORef)
import GHC.STRef (STRef (..))
import GHC.Exts ( casMutVar#, MutVar#, RealWorld, readMutVar#, lazy, State#
, writeMutVar#, Int (..), Int#, MutableArray#, readArray#
, casArray#, SmallMutableArray#, readSmallArray#, casSmallArray#
)
import Data.IORef.AtomicModify.SmallArraySize
( getSizeofSmallMutableArray# )
import GHC.IO (IO (..))
import System.IO.Unsafe (unsafeDupablePerformIO)
import Data.Primitive.Array
import Data.Primitive.SmallArray
import Control.Monad (when)
import Control.Exception (ArrayException (..), throwIO)
-- | A version of 'GHC.IORef.atomicModifyIORef2' that takes an arbitrary pair
-- of functions. This function will allocate more than 'atomicModifyIORef2',
-- and will tend to take longer to succeed when there is a lot of contention
-- for the 'IORef'.
--
-- @
-- atomicModifyIORef2 ref f = do
-- (old, _new, r) <- atomicModifyIORef2General ref fst f
-- pure (old, r)
-- @
--
-- If the first function (the \"extraction function\") is a record field
-- selector (e.g., 'snd'), we do our best to make sure the thunk placed in the
-- 'IORef' is a selector thunk, so the garbage collector can drop the rest of
-- the record once the record is forced. In other cases, callers should
-- generally force the returned @new@ value in order to avoid a potential space
-- leak.
--
-- Conceptually:
--
-- @
-- atomicModifyIORef3General ref extract f = do
-- -- Begin atomic block
-- old <- 'readIORef' ref
-- let r = f old
-- new = extract r
-- 'writeIORef' ref new
-- -- End atomic block
-- r `seq` pure (old, new, r)
-- @
--
-- where other threads cannot interfere with the operations in the \"atomic block\".
-- In particular, no other thread can write to the 'IORef' between the 'readIORef'
-- and the 'writeIORef' operations.
atomicModifyIORef3General
:: IORef a -> (t -> a) -> (a -> t) -> IO (a, a, t)
atomicModifyIORef3General (IORef (STRef ref)) extract = \f -> do
-- atomicModifyMutVar2# creates a thunk for the result of applying the user
-- function 'f' to the "old" value read from the IORef, and then edits that thunk
-- in a CAS loop. In Haskell land, thunks are immutable, so we can't exactly
-- do that. Instead, we make an IORef, 'holder', to hold the "old" value, and
-- use 'unsafeDupablePerformIO' to create a thunk that will read it and apply
-- 'f' to the result. We /can/ edit the holder IORef in the CAS loop.
-- Note: since casMutVar# introduces a full memory barrier, any thread reading
-- the 'r' thunk from 'ref' will have "seen" the preceding 'writeMutVar' to 'holder',
-- so it won't get the uninitialized value or anything similarly stale.
holder@(IORef (STRef holder#)) <- newIORef uninitialized
let r = unsafeDupablePerformIO (f <$> readIORef holder)
-- I really don't think r is going to inline anyway, but if it does, we
-- could produce an unnecessary space leak.
{-# NOINLINE r #-}
let new = extract r
IO (\s -> case atomicModifyIORef3General' ref holder# new r s of
(# s', old, new', !res #) -> (# s', (old, new', res) #))
{-# INLINE atomicModifyIORef3General #-}
atomicModifyIORef3General'
:: MutVar# RealWorld a -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifyIORef3General' ref holder new r s1 =
case readMutVar# ref s1 of { (# s2, old #) ->
case writeMutVar# holder old s2 of { s3 ->
case casMutVar# ref old new s3 of { (# s4, flag, _ #) ->
case flag of
-- Why the lazy invocations?
--
-- In the event that 'old' gets forced, unboxed, and reboxed between
-- 'readMutVar#' and 'casMutVar#', the CAS will never succeed. I doubt
-- that could happen anyway, but let's be sure.
--
-- If 'r' is forced before the holder is written, it will read an
-- uninitialized value and throw an exception. Ouch. Let's make sure that
-- doesn't happen either.
0# -> (# s4, lazy old, lazy new, lazy r #)
_ -> atomicModifyIORef3General' ref holder new r s4 }}}
{-# NOINLINE atomicModifyIORef3General' #-}
uninitialized :: a
uninitialized = error "Uninitialized. This is a bug in atomic-modify-generics."
{-# NOINLINE uninitialized #-}
-- | A version of 'atomicModifyIORef3General' for 'Array's. See the
-- documentation there. Indexing is performed safely.
atomicModifyArray3General
:: MutableArray RealWorld a -> Int -> (t -> a) -> (a -> t) -> IO (a, a, t)
-- See atomicModifyIORef3General for implementation comments.
-- Why do we perform safe indexing? This operation is expensive enough
-- that I don't think we really have to worry about the cost of a single
-- bounds check.
atomicModifyArray3General mary@(MutableArray mary#) ix@(I# ix#) extract = \f -> do
-- We use unsigned comparison to make sure ix is non-negative
-- and less than the array size with just one comparison. The LLVM
-- backend is clever enough to produce this from the obvious two-sided
-- check, but last I looked the native code generator wasn't.
let !sz = sizeofMutableArray mary
when ((fromIntegral ix :: Word) >= fromIntegral sz) $ outOfBoundsArray ix sz
holder@(IORef (STRef holder#)) <- newIORef uninitialized
let r = unsafeDupablePerformIO (f <$> readIORef holder)
{-# NOINLINE r #-}
let new = extract r
IO (\s -> case atomicModifyArray3General' mary# ix# holder# new r s of
(# s', old, new', !res #) -> (# s', (old, new', res) #))
{-# INLINE atomicModifyArray3General #-}
outOfBoundsArray :: Int -> Int -> IO a
outOfBoundsArray ix sz
| ix < 0 = throwIO . IndexOutOfBounds $
"\natomicModifyArray3General was passed a negative array index of " ++
show ix ++ "."
| otherwise = throwIO . IndexOutOfBounds $
"\natomicModifyArray3General was passed an array index of " ++
show ix ++ ",\nbut an array of only " ++ show sz ++ " elements."
{-# NOINLINE outOfBoundsArray #-}
atomicModifyArray3General'
:: MutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifyArray3General' mary ix holder new r s1 =
case readArray# mary ix s1 of { (# s2, old #) ->
case writeMutVar# holder old s2 of { s3 ->
case casArray# mary ix old new s3 of { (# s4, flag, _ #) ->
case flag of
0# -> (# s4, lazy old, lazy new, lazy r #)
_ -> atomicModifyArray3General' mary ix holder new r s4 }}}
{-# NOINLINE atomicModifyArray3General' #-}
-- | A version of 'atomicModifyIORef3General' for 'SmallArray's. See the
-- documentation there. Indexing is performed safely.
atomicModifySmallArray3General
:: SmallMutableArray RealWorld a -> Int -> (t -> a) -> (a -> t) -> IO (a, a, t)
-- See atomicModifyIORef3General for implementation comments.
atomicModifySmallArray3General (SmallMutableArray mary#) ix@(I# ix#) extract = \f -> do
IO (\s -> case getSizeofSmallMutableArray# mary# s of (# s', sz# #) -> (# s', I# sz# #))
>>= \sz -> when ((fromIntegral ix :: Word) >= fromIntegral sz) (outOfBoundsSmallArray ix sz)
holder@(IORef (STRef holder#)) <- newIORef uninitialized
let r = unsafeDupablePerformIO (f <$> readIORef holder)
{-# NOINLINE r #-}
let new = extract r
IO (\s -> case atomicModifySmallArray3General' mary# ix# holder# new r s of
(# s', old, new', !res #) -> (# s', (old, new', res) #))
{-# INLINE atomicModifySmallArray3General #-}
atomicModifySmallArray3General'
:: SmallMutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifySmallArray3General' mary ix holder new r s1 =
case readSmallArray# mary ix s1 of { (# s2, old #) ->
case writeMutVar# holder old s2 of { s3 ->
case casSmallArray# mary ix old new s3 of { (# s4, flag, _ #) ->
case flag of
0# -> (# s4, lazy old, lazy new, lazy r #)
_ -> atomicModifySmallArray3General' mary ix holder new r s4 }}}
{-# NOINLINE atomicModifySmallArray3General' #-}
outOfBoundsSmallArray :: Int -> Int -> IO a
outOfBoundsSmallArray ix sz
| ix < 0 = throwIO . IndexOutOfBounds $
"\natomicModifySmallArray3General was passed a negative array index of " ++
show ix ++ "."
| otherwise = throwIO . IndexOutOfBounds $
"\natomicModifySmallArray3General was passed an array index of " ++
show ix ++ ",\nbut an array of only " ++ show sz ++ " elements."
{-# NOINLINE outOfBoundsSmallArray #-}