{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

{- |
Module      :  OpenTelemetry.Internal.AtomicCounter
Copyright   :  (c) Ian Duncan, 2024-2026
License     :  BSD-3
Description :  Machine-word atomic counter using hardware fetch-and-add.
Stability   :  experimental

Unlike @atomicModifyIORef'@ which does a CAS retry loop on the boxed
@IORef@ closure, these operations compile down to a single
@lock xadd@ (x86) or @ldadd@ (AArch64) instruction with no
allocation and no retry.
-}
module OpenTelemetry.Internal.AtomicCounter (
  AtomicCounter,
  newAtomicCounter,
  incrAtomicCounter,
  addAtomicCounter,
  fetchAddAtomicCounter,
  readAtomicCounter,
  writeAtomicCounter,
) where


#include "MachDeps.h"

import GHC.Exts (
  Int (..),
  MutableByteArray#,
  RealWorld,
  fetchAddIntArray#,
  newByteArray#,
  readIntArray#,
  writeIntArray#,
  (+#),
 )
import GHC.IO (IO (..))


{- | A mutable atomic counter backed by a single unboxed machine-word 'Int'.

Uses hardware fetch-and-add (@fetchAddIntArray#@) instead of CAS retry
loops, making increment\/add O(1) regardless of contention.
-}
data AtomicCounter = AtomicCounter (MutableByteArray# RealWorld)


-- | Create a new counter initialized to the given value.
newAtomicCounter :: Int -> IO AtomicCounter
newAtomicCounter :: Int -> IO AtomicCounter
newAtomicCounter (I# Int#
n) = (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# SIZEOF_HSINT# s of
    (# State# RealWorld
s1, MutableByteArray# RealWorld
arr #) -> 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#
n State# RealWorld
s1 of
      State# RealWorld
s2 -> (# State# RealWorld
s2, MutableByteArray# RealWorld -> AtomicCounter
AtomicCounter MutableByteArray# RealWorld
arr #)


-- | Atomically increment the counter by 1. Returns the value /after/ the increment.
incrAtomicCounter :: AtomicCounter -> IO Int
incrAtomicCounter :: AtomicCounter -> IO Int
incrAtomicCounter = Int -> AtomicCounter -> IO Int
addAtomicCounter Int
1
{-# INLINE incrAtomicCounter #-}


-- | Atomically add to the counter. Returns the value /after/ the add.
addAtomicCounter :: Int -> AtomicCounter -> IO Int
addAtomicCounter :: Int -> AtomicCounter -> IO Int
addAtomicCounter (I# Int#
incr) (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# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
arr Int#
0# Int#
incr State# RealWorld
s of
    (# State# RealWorld
s', Int#
old #) -> (# State# RealWorld
s', Int# -> Int
I# (Int#
old Int# -> Int# -> Int#
+# Int#
incr) #)
{-# INLINE addAtomicCounter #-}


{- | Atomically add to the counter. Returns the value /before/ the add.
Useful for monotonic ID allocation.
-}
fetchAddAtomicCounter :: Int -> AtomicCounter -> IO Int
fetchAddAtomicCounter :: Int -> AtomicCounter -> IO Int
fetchAddAtomicCounter (I# Int#
incr) (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# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
arr Int#
0# Int#
incr State# RealWorld
s of
    (# State# RealWorld
s', Int#
old #) -> (# State# RealWorld
s', Int# -> Int
I# Int#
old #)
{-# INLINE fetchAddAtomicCounter #-}


{- | Read the current counter value.

This is a relaxed read; no ordering guarantee relative to concurrent adds
on other cores. Fine for diagnostics and metrics.
-}
readAtomicCounter :: AtomicCounter -> IO Int
readAtomicCounter :: AtomicCounter -> IO Int
readAtomicCounter (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#
val #) -> (# State# RealWorld
s', Int# -> Int
I# Int#
val #)
{-# INLINE readAtomicCounter #-}


-- | Overwrite the counter value. Not atomic with respect to concurrent adds.
writeAtomicCounter :: AtomicCounter -> Int -> IO ()
writeAtomicCounter :: AtomicCounter -> Int -> IO ()
writeAtomicCounter (AtomicCounter MutableByteArray# RealWorld
arr) (I# Int#
n) = (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#
n State# RealWorld
s of
    State# RealWorld
s' -> (# State# RealWorld
s', () #)
{-# INLINE writeAtomicCounter #-}