{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- | This module provides a simple, efficient supply of integers using atomic fetch-and-add.
--
-- To use this module, first create an @IntSupply@. This is often done once at the top level of an application, in
-- global scope.
--
-- > import IntSupply (IntSupply)
-- > import IntSupply qualified
-- > import System.IO.Unsafe (unsafePerformIO)
-- >
-- > myIntSupply :: IntSupply
-- > myIntSupply = unsafePerformIO IntSupply.new
-- > {-# NOINLINE myIntSupply #-}
--
-- Next, call @IntSupply.next@ on the supply, which will return 0, then 1, and so on.
--
-- > > IntSupply.next myIntSupply
-- > 0
-- > > IntSupply.next myIntSupply
-- > 1
--
-- If desired, you can reset the count to 0.
--
-- > > IntSupply.reset myIntSupply
-- > > IntSupply.next myIntSupply
-- > 0
--
-- On a 64-bit machine, for many applications, these integers can be treated as effectively unique: even if
-- 1,000,000,000 integers were generated per second, it would still take over 580 years to wrap around.
--
-- On a 32-bit machine, more care must be taken, of course: even if only 1,000 integers were generated per second, it
-- would only take 50 days to wrap around.
module IntSupply
  ( IntSupply,
    new,
    next,
    reset,
  )
where

import Data.Bits (finiteBitSize)
import GHC.Base
  ( IO (IO),
    Int (I#),
    MutableByteArray#,
    RealWorld,
    atomicWriteIntArray#,
    fetchAddIntArray#,
    newByteArray#,
    writeIntArray#,
  )

-- | A thread-safe supply of integers.
data IntSupply
  = IntSupply (MutableByteArray# RealWorld)

-- | Create a supply of integers.
new :: IO IntSupply
new :: IO IntSupply
new =
  (State# RealWorld -> (# State# RealWorld, IntSupply #))
-> IO IntSupply
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s0 ->
    case Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
size State# RealWorld
s0 of
      (# State# RealWorld
s1, MutableByteArray# RealWorld
supply #) ->
        (# MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
writeIntArray# MutableByteArray# RealWorld
supply Int#
0# Int#
0# State# RealWorld
s1, MutableByteArray# RealWorld -> IntSupply
IntSupply MutableByteArray# RealWorld
supply #)
  where
    !(I# Int#
size) =
      Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
forall a. HasCallStack => a
undefined :: Int) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8
{-# INLINEABLE new #-}

-- | Get the next integer from a supply of integers.
next :: IntSupply -> IO Int
next :: IntSupply -> IO Int
next (IntSupply MutableByteArray# RealWorld
supply) =
  (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s0 ->
    case MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
supply Int#
0# Int#
1# State# RealWorld
s0 of
      (# State# RealWorld
s1, Int#
n #) -> (# State# RealWorld
s1, Int# -> Int
I# Int#
n #)
{-# INLINEABLE next #-}

-- | Reset a supply of integers to 0.
reset :: IntSupply -> IO ()
reset :: IntSupply -> IO ()
reset (IntSupply MutableByteArray# RealWorld
arr#) =
  (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO \State# RealWorld
s0 ->
    (# MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
atomicWriteIntArray# MutableByteArray# RealWorld
arr# Int#
0# Int#
0# State# RealWorld
s0, () #)
{-# INLINEABLE reset #-}