{-# 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 :: forall a t. IORef a -> (t -> a) -> (a -> t) -> IO (a, a, t)
atomicModifyIORef3General (IORef (STRef MutVar# RealWorld a
ref)) t -> a
extract = \a -> t
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 a
holder@(IORef (STRef MutVar# RealWorld a
holder#)) <- forall a. a -> IO (IORef a)
newIORef forall a. a
uninitialized
  let r :: t
r = forall a. IO a -> a
unsafeDupablePerformIO (a -> t
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef a
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 :: a
new = t -> a
extract t
r
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> case forall a t.
MutVar# RealWorld a
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyIORef3General' MutVar# RealWorld a
ref MutVar# RealWorld a
holder# a
new t
r State# RealWorld
s of
               (# State# RealWorld
s', a
old, a
new', !t
res #) -> (# State# RealWorld
s', (a
old, a
new', t
res) #))
{-# INLINE atomicModifyIORef3General #-}

atomicModifyIORef3General'
  :: MutVar# RealWorld a -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifyIORef3General' :: forall a t.
MutVar# RealWorld a
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyIORef3General' MutVar# RealWorld a
ref MutVar# RealWorld a
holder a
new t
r State# RealWorld
s1 =
  case forall d a. MutVar# d a -> State# d -> (# State# d, a #)
readMutVar# MutVar# RealWorld a
ref State# RealWorld
s1 of { (# State# RealWorld
s2, a
old #) ->
  case forall d a. MutVar# d a -> a -> State# d -> State# d
writeMutVar# MutVar# RealWorld a
holder a
old State# RealWorld
s2 of { State# RealWorld
s3 ->
  case forall d a.
MutVar# d a -> a -> a -> State# d -> (# State# d, Int#, a #)
casMutVar# MutVar# RealWorld a
ref a
old a
new State# RealWorld
s3 of { (# State# RealWorld
s4, Int#
flag, a
_ #) ->
  case Int#
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.
    Int#
0# -> (# State# RealWorld
s4, forall a. a -> a
lazy a
old, forall a. a -> a
lazy a
new, forall a. a -> a
lazy t
r #)
    Int#
_ -> forall a t.
MutVar# RealWorld a
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyIORef3General' MutVar# RealWorld a
ref MutVar# RealWorld a
holder a
new t
r State# RealWorld
s4 }}}
{-# NOINLINE atomicModifyIORef3General' #-}

uninitialized :: a
uninitialized :: forall a. a
uninitialized = forall a. HasCallStack => [Char] -> a
error [Char]
"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 :: forall a t.
MutableArray RealWorld a
-> Int -> (t -> a) -> (a -> t) -> IO (a, a, t)
atomicModifyArray3General mary :: MutableArray RealWorld a
mary@(MutableArray MutableArray# RealWorld a
mary#) ix :: Int
ix@(I# Int#
ix#) t -> a
extract = \a -> t
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 :: Int
sz = forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray RealWorld a
mary
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz) forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> IO a
outOfBoundsArray Int
ix Int
sz
  holder :: IORef a
holder@(IORef (STRef MutVar# RealWorld a
holder#)) <- forall a. a -> IO (IORef a)
newIORef forall a. a
uninitialized
  let r :: t
r = forall a. IO a -> a
unsafeDupablePerformIO (a -> t
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef a
holder)
      {-# NOINLINE r #-}
  let new :: a
new = t -> a
extract t
r
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> case forall a t.
MutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyArray3General' MutableArray# RealWorld a
mary# Int#
ix# MutVar# RealWorld a
holder# a
new t
r State# RealWorld
s of
               (# State# RealWorld
s', a
old, a
new', !t
res #) -> (# State# RealWorld
s', (a
old, a
new', t
res) #))
{-# INLINE atomicModifyArray3General #-}

outOfBoundsArray :: Int -> Int -> IO a
outOfBoundsArray :: forall a. Int -> Int -> IO a
outOfBoundsArray Int
ix Int
sz
  | Int
ix forall a. Ord a => a -> a -> Bool
< Int
0 = forall e a. Exception e => e -> IO a
throwIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ArrayException
IndexOutOfBounds forall a b. (a -> b) -> a -> b
$
      [Char]
"\natomicModifyArray3General was passed a negative array index of " forall a. [a] -> [a] -> [a]
++
        forall a. Show a => a -> [Char]
show Int
ix forall a. [a] -> [a] -> [a]
++ [Char]
"."
  | Bool
otherwise = forall e a. Exception e => e -> IO a
throwIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ArrayException
IndexOutOfBounds forall a b. (a -> b) -> a -> b
$
      [Char]
"\natomicModifyArray3General was passed an array index of " forall a. [a] -> [a] -> [a]
++
        forall a. Show a => a -> [Char]
show Int
ix forall a. [a] -> [a] -> [a]
++ [Char]
",\nbut an array of only " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
sz forall a. [a] -> [a] -> [a]
++ [Char]
" elements."
{-# NOINLINE outOfBoundsArray #-}

atomicModifyArray3General'
  :: MutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifyArray3General' :: forall a t.
MutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyArray3General' MutableArray# RealWorld a
mary Int#
ix MutVar# RealWorld a
holder a
new t
r State# RealWorld
s1 =
  case forall d a.
MutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# MutableArray# RealWorld a
mary Int#
ix State# RealWorld
s1 of { (# State# RealWorld
s2, a
old #) ->
  case forall d a. MutVar# d a -> a -> State# d -> State# d
writeMutVar# MutVar# RealWorld a
holder a
old State# RealWorld
s2 of { State# RealWorld
s3 ->
  case forall d a.
MutableArray# d a
-> Int# -> a -> a -> State# d -> (# State# d, Int#, a #)
casArray# MutableArray# RealWorld a
mary Int#
ix a
old a
new State# RealWorld
s3 of { (# State# RealWorld
s4, Int#
flag, a
_ #) ->
  case Int#
flag of
    Int#
0# -> (# State# RealWorld
s4, forall a. a -> a
lazy a
old, forall a. a -> a
lazy a
new, forall a. a -> a
lazy t
r #)
    Int#
_ -> forall a t.
MutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifyArray3General' MutableArray# RealWorld a
mary Int#
ix MutVar# RealWorld a
holder a
new t
r State# RealWorld
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 :: forall a t.
SmallMutableArray RealWorld a
-> Int -> (t -> a) -> (a -> t) -> IO (a, a, t)
atomicModifySmallArray3General (SmallMutableArray SmallMutableArray# RealWorld a
mary#) ix :: Int
ix@(I# Int#
ix#) t -> a
extract = \a -> t
f -> do
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> case forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, Int# #)
getSizeofSmallMutableArray# SmallMutableArray# RealWorld a
mary# State# RealWorld
s of (# State# RealWorld
s', Int#
sz# #) -> (# State# RealWorld
s', Int# -> Int
I# Int#
sz# #))
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
sz -> forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz) (forall a. Int -> Int -> IO a
outOfBoundsSmallArray Int
ix Int
sz)
  holder :: IORef a
holder@(IORef (STRef MutVar# RealWorld a
holder#)) <- forall a. a -> IO (IORef a)
newIORef forall a. a
uninitialized
  let r :: t
r = forall a. IO a -> a
unsafeDupablePerformIO (a -> t
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef a
holder)
      {-# NOINLINE r #-}
  let new :: a
new = t -> a
extract t
r
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> case forall a t.
SmallMutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifySmallArray3General' SmallMutableArray# RealWorld a
mary# Int#
ix# MutVar# RealWorld a
holder# a
new t
r State# RealWorld
s of
               (# State# RealWorld
s', a
old, a
new', !t
res #) -> (# State# RealWorld
s', (a
old, a
new', t
res) #))
{-# INLINE atomicModifySmallArray3General #-}

atomicModifySmallArray3General'
  :: SmallMutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #)
atomicModifySmallArray3General' :: forall a t.
SmallMutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifySmallArray3General' SmallMutableArray# RealWorld a
mary Int#
ix MutVar# RealWorld a
holder a
new t
r State# RealWorld
s1 =
  case forall d a.
SmallMutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readSmallArray# SmallMutableArray# RealWorld a
mary Int#
ix State# RealWorld
s1 of { (# State# RealWorld
s2, a
old #) ->
  case forall d a. MutVar# d a -> a -> State# d -> State# d
writeMutVar# MutVar# RealWorld a
holder a
old State# RealWorld
s2 of { State# RealWorld
s3 ->
  case forall d a.
SmallMutableArray# d a
-> Int# -> a -> a -> State# d -> (# State# d, Int#, a #)
casSmallArray# SmallMutableArray# RealWorld a
mary Int#
ix a
old a
new State# RealWorld
s3 of { (# State# RealWorld
s4, Int#
flag, a
_ #) ->
  case Int#
flag of
    Int#
0# -> (# State# RealWorld
s4, forall a. a -> a
lazy a
old, forall a. a -> a
lazy a
new, forall a. a -> a
lazy t
r #)
    Int#
_ -> forall a t.
SmallMutableArray# RealWorld a
-> Int#
-> MutVar# RealWorld a
-> a
-> t
-> State# RealWorld
-> (# State# RealWorld, a, a, t #)
atomicModifySmallArray3General' SmallMutableArray# RealWorld a
mary Int#
ix MutVar# RealWorld a
holder a
new t
r State# RealWorld
s4 }}}
{-# NOINLINE atomicModifySmallArray3General' #-}

outOfBoundsSmallArray :: Int -> Int -> IO a
outOfBoundsSmallArray :: forall a. Int -> Int -> IO a
outOfBoundsSmallArray Int
ix Int
sz
  | Int
ix forall a. Ord a => a -> a -> Bool
< Int
0 = forall e a. Exception e => e -> IO a
throwIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ArrayException
IndexOutOfBounds forall a b. (a -> b) -> a -> b
$
      [Char]
"\natomicModifySmallArray3General was passed a negative array index of " forall a. [a] -> [a] -> [a]
++
        forall a. Show a => a -> [Char]
show Int
ix forall a. [a] -> [a] -> [a]
++ [Char]
"."
  | Bool
otherwise = forall e a. Exception e => e -> IO a
throwIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ArrayException
IndexOutOfBounds forall a b. (a -> b) -> a -> b
$
      [Char]
"\natomicModifySmallArray3General was passed an array index of " forall a. [a] -> [a] -> [a]
++
        forall a. Show a => a -> [Char]
show Int
ix forall a. [a] -> [a] -> [a]
++ [Char]
",\nbut an array of only " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
sz forall a. [a] -> [a] -> [a]
++ [Char]
" elements."
{-# NOINLINE outOfBoundsSmallArray #-}