{-# LANGUAGE BangPatterns, MagicHash, UnboxedTuples, CPP #-}
-- | Integer counters providing thread-safe, lock-free mutation functions.
--
--   Atomic counters are represented by a single memory location, such that
--   built-in processor instructions are sufficient to perform fetch-and-add or
--   compare-and-swap.
-- 
--   Remember, contention on such counters should still be minimized!

module Data.Atomics.Counter
       -- Reexport to get all the docs.
       (
         -- * Type of counters of counters and tickets
         AtomicCounter, 
         
         -- * Creating counters
         newCounter, 

         -- * Tickets, used for compare-and-swap         
         -- | See the documentation for "Data.Atomics" for more explanation of the
         -- ticket abstraction.  The same ideas apply here for counters as for
         -- general mutable locations (IORefs).
         CTicket, peekCTicket,

         -- * Atomic memory operations
         casCounter, incrCounter, incrCounter_,
                                  
         -- * Non-atomic operations
         readCounter, readCounterForCAS,
         writeCounter
         )
 where


import Data.Atomics.Internal
import GHC.Base  hiding ((==#))
import qualified GHC.PrimopWrappers as GPW


-- GHC 7.8 changed some primops
(==#) :: Int# -> Int# -> Bool
==# :: Int# -> Int# -> Bool
(==#) Int#
x Int#
y = case Int#
x Int# -> Int# -> Int#
GPW.==# Int#
y of { Int#
0# -> Bool
False; Int#
_ -> Bool
True }



#ifndef __GLASGOW_HASKELL__
#error "Counter: this library is not portable to other Haskell's"
#endif

#include "MachDeps.h"
#ifndef SIZEOF_HSINT
#define SIZEOF_HSINT  INT_SIZE_IN_BYTES
#endif

-- | The type of mutable atomic counters.
data AtomicCounter = AtomicCounter (MutableByteArray# RealWorld)

-- | You should not depend on this type.  It varies between different implementations
-- of atomic counters.
type CTicket = Int
-- TODO: Could newtype this.

-- | Create a new counter initialized to the given value.
{-# INLINE newCounter #-}
newCounter :: Int -> IO AtomicCounter
newCounter :: Int -> IO AtomicCounter
newCounter Int
n = do
  AtomicCounter
c <- IO AtomicCounter
newRawCounter
  AtomicCounter -> Int -> IO ()
writeCounter AtomicCounter
c Int
n -- Non-atomic is ok; it hasn't been released into the wild.
  AtomicCounter -> IO AtomicCounter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return AtomicCounter
c

-- | Create a new, uninitialized counter.
{-# INLINE newRawCounter #-}
newRawCounter :: IO AtomicCounter  
newRawCounter :: IO AtomicCounter
newRawCounter = (State# RealWorld -> (# State# RealWorld, AtomicCounter #))
-> IO AtomicCounter
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, AtomicCounter #))
 -> IO AtomicCounter)
-> (State# RealWorld -> (# State# RealWorld, AtomicCounter #))
-> IO AtomicCounter
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
size State# RealWorld
s of { (# State# RealWorld
s', MutableByteArray# RealWorld
arr #) ->
  (# State# RealWorld
s', MutableByteArray# RealWorld -> AtomicCounter
AtomicCounter MutableByteArray# RealWorld
arr #) }
  where !(I# Int#
size) = SIZEOF_HSINT

{-# INLINE readCounter #-}
-- | Equivalent to `readCounterForCAS` followed by `peekCTicket`.        
readCounter :: AtomicCounter -> IO Int
readCounter :: AtomicCounter -> IO Int
readCounter (AtomicCounter MutableByteArray# RealWorld
arr) = (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int)
-> (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
readIntArray# MutableByteArray# RealWorld
arr Int#
0# State# RealWorld
s of { (# State# RealWorld
s', Int#
i #) ->
  (# State# RealWorld
s', Int# -> Int
I# Int#
i #) }

{-# INLINE writeCounter #-}
-- | Make a non-atomic write to the counter.  No memory-barrier.
writeCounter :: AtomicCounter -> Int -> IO ()
writeCounter :: AtomicCounter -> Int -> IO ()
writeCounter (AtomicCounter MutableByteArray# RealWorld
arr) (I# Int#
i) = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
writeIntArray# MutableByteArray# RealWorld
arr Int#
0# Int#
i State# RealWorld
s of { State# RealWorld
s' ->
  (# State# RealWorld
s', () #) }

{-# INLINE readCounterForCAS #-}
-- | Just like the "Data.Atomics" CAS interface, this routine returns an opaque
-- ticket that can be used in CAS operations.  Except for the difference in return
-- type, the semantics of this are the same as `readCounter`.
readCounterForCAS :: AtomicCounter -> IO CTicket
readCounterForCAS :: AtomicCounter -> IO Int
readCounterForCAS = AtomicCounter -> IO Int
readCounter

{-# INLINE peekCTicket #-}
-- | Opaque tickets cannot be constructed, but they can be destructed into values.
peekCTicket :: CTicket -> Int
peekCTicket :: Int -> Int
peekCTicket !Int
x = Int
x

{-# INLINE casCounter #-}
-- | Compare and swap for the counter ADT.  Similar behavior to
-- `Data.Atomics.casIORef`, in particular, in both success and failure cases it
-- returns a ticket that you should use for the next attempt.  (That is, in the
-- success case, it actually returns the new value that you provided as input, but in
-- ticket form.)
casCounter :: AtomicCounter -> CTicket -> Int -> IO (Bool, CTicket)
-- casCounter (AtomicCounter barr) !old !new =
casCounter :: AtomicCounter -> Int -> Int -> IO (Bool, Int)
casCounter (AtomicCounter MutableByteArray# RealWorld
mba#) (I# Int#
old#) newBox :: Int
newBox@(I# Int#
new#) = (State# RealWorld -> (# State# RealWorld, (Bool, Int) #))
-> IO (Bool, Int)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO((State# RealWorld -> (# State# RealWorld, (Bool, Int) #))
 -> IO (Bool, Int))
-> (State# RealWorld -> (# State# RealWorld, (Bool, Int) #))
-> IO (Bool, Int)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s1# ->
  let (# State# RealWorld
s2#, Int#
res# #) = MutableByteArray# RealWorld
-> Int#
-> Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
casIntArray# MutableByteArray# RealWorld
mba# Int#
0# Int#
old# Int#
new# State# RealWorld
s1# in
  case Int#
res# Int# -> Int# -> Bool
==# Int#
old# of 
    Bool
False -> (# State# RealWorld
s2#, (Bool
False, Int# -> Int
I# Int#
res# ) #) -- Failure
    Bool
True  -> (# State# RealWorld
s2#, (Bool
True , Int
newBox ) #) -- Success

-- {-# INLINE sameCTicket #-}
-- sameCTicket :: CTicket -> CTicket -> Bool
-- sameCTicket = (==)

{-# INLINE incrCounter #-}
-- | Increment the counter by a given amount.  Returns the value AFTER the increment
--   (in contrast with the behavior of the underlying instruction on architectures
--   like x86.)
--                 
--   Note that UNLIKE with boxed implementations of counters, where increment is
--   based on CAS, this increment is /O(1)/.  Fetch-and-add does not require a retry
--   loop like CAS.
incrCounter :: Int -> AtomicCounter -> IO Int
incrCounter :: Int -> AtomicCounter -> IO Int
incrCounter (I# Int#
incr#) (AtomicCounter MutableByteArray# RealWorld
mba#) = (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int)
-> (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \ State# RealWorld
s1# -> 
  let (# State# RealWorld
s2#, Int#
res #) = MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
mba# Int#
0# Int#
incr# State# RealWorld
s1# in
  (# State# RealWorld
s2#, (Int# -> Int
I# (Int#
res Int# -> Int# -> Int#
+# Int#
incr#)) #)

{-# INLINE incrCounter_ #-}
-- | An alternate version for when you don't care about the old value.
incrCounter_ :: Int -> AtomicCounter -> IO ()
incrCounter_ :: Int -> AtomicCounter -> IO ()
incrCounter_ (I# Int#
incr#) (AtomicCounter MutableByteArray# RealWorld
mba#) = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ State# RealWorld
s1# -> 
  let (# State# RealWorld
s2#, Int#
_ #) = MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
mba# Int#
0# Int#
incr# State# RealWorld
s1# in
  (# State# RealWorld
s2#, () #)