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

{- |
Module      :  OpenTelemetry.Internal.AtomicBucketArray
Copyright   :  (c) Ian Duncan, 2024-2026
License     :  BSD-3
Description :  Contiguous array of atomic counters for histogram bucket counts.
Stability   :  experimental

One 'MutableByteArray#' holds all bucket counts in a single allocation,
indexed by bucket number. Each slot supports atomic fetch-and-add via
@fetchAddIntArray#@, a single @ldadd@ (AArch64) or @lock xadd@ (x86)
instruction with no CAS retry and no allocation.

Used to eliminate the O(n) bucket vector copy that 'Data.Vector.Unboxed.modify'
performs on every histogram recording inside an @atomicModifyIORef'@ CAS loop.
-}
module OpenTelemetry.Internal.AtomicBucketArray (
  AtomicBucketArray,
  newAtomicBucketArray,
  atomicAddBucket,
  readBucketArray,
  readAndResetBucketArray,
) where


#include "MachDeps.h"

import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as U
import Data.Word (Word64)
import GHC.Exts (
  Int (..),
  Int#,
  MutableByteArray#,
  RealWorld,
  State#,
  fetchAddIntArray#,
  isTrue#,
  negateInt#,
  newByteArray#,
  readIntArray#,
  writeIntArray#,
  (*#),
  (+#),
  (>=#),
 )
import GHC.IO (IO (..))


{- | Contiguous mutable array of atomic counters. Each element is one
machine-word 'Int', accessible via hardware fetch-and-add.

@since 0.0.1.0
-}
data AtomicBucketArray = AtomicBucketArray (MutableByteArray# RealWorld) Int#


{- | Allocate a new bucket array with all counters initialized to zero.

@since 0.0.1.0
-}
newAtomicBucketArray :: Int -> IO AtomicBucketArray
newAtomicBucketArray :: Int -> IO AtomicBucketArray
newAtomicBucketArray (I# Int#
n) = (State# RealWorld -> (# State# RealWorld, AtomicBucketArray #))
-> IO AtomicBucketArray
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, AtomicBucketArray #))
 -> IO AtomicBucketArray)
-> (State# RealWorld -> (# State# RealWorld, AtomicBucketArray #))
-> IO AtomicBucketArray
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  let !nbytes :: Int#
nbytes = Int#
n Int# -> Int# -> Int#
*# SIZEOF_HSINT#
  in case Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
nbytes State# RealWorld
s of
       (# State# RealWorld
s1, MutableByteArray# RealWorld
arr #) -> case MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> State# RealWorld
zeroFill MutableByteArray# RealWorld
arr Int#
n State# RealWorld
s1 of
         State# RealWorld
s2 -> (# State# RealWorld
s2, MutableByteArray# RealWorld -> Int# -> AtomicBucketArray
AtomicBucketArray MutableByteArray# RealWorld
arr Int#
n #)


zeroFill :: MutableByteArray# RealWorld -> Int# -> State# RealWorld -> State# RealWorld
zeroFill :: MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> State# RealWorld
zeroFill MutableByteArray# RealWorld
arr Int#
nwords State# RealWorld
s0 = Int# -> State# RealWorld -> State# RealWorld
go Int#
0# State# RealWorld
s0
  where
    go :: Int# -> State# RealWorld -> State# RealWorld
go Int#
i State# RealWorld
s
      | Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
>=# Int#
nwords) = State# RealWorld
s
      | Bool
otherwise = case MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
writeIntArray# MutableByteArray# RealWorld
arr Int#
i Int#
0# State# RealWorld
s of
          State# RealWorld
s' -> Int# -> State# RealWorld -> State# RealWorld
go (Int#
i Int# -> Int# -> Int#
+# Int#
1#) State# RealWorld
s'


{- | Atomically increment the bucket at the given index by 1.
No bounds checking.

@since 0.0.1.0
-}
atomicAddBucket :: AtomicBucketArray -> Int -> IO ()
atomicAddBucket :: AtomicBucketArray -> Int -> IO ()
atomicAddBucket (AtomicBucketArray MutableByteArray# RealWorld
arr Int#
_) (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, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
arr Int#
i Int#
1# State# RealWorld
s of
    (# State# RealWorld
s', Int#
_ #) -> (# State# RealWorld
s', () #)
{-# INLINE atomicAddBucket #-}


{- | Read all bucket values into an immutable 'U.Vector Word64'.
Each element is read individually; not globally atomic across buckets
(acceptable for metric snapshots).

@since 0.0.1.0
-}
readBucketArray :: AtomicBucketArray -> IO (Vector Word64)
readBucketArray :: AtomicBucketArray -> IO (Vector Word64)
readBucketArray (AtomicBucketArray MutableByteArray# RealWorld
arr Int#
len) = do
  let !n :: Int
n = Int# -> Int
I# Int#
len
  Int -> (Int -> IO Word64) -> IO (Vector Word64)
forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> (Int -> m a) -> m (Vector a)
U.generateM Int
n ((Int -> IO Word64) -> IO (Vector Word64))
-> (Int -> IO Word64) -> IO (Vector Word64)
forall a b. (a -> b) -> a -> b
$ \(I# Int#
i) -> (State# RealWorld -> (# State# RealWorld, Word64 #)) -> IO Word64
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Word64 #)) -> IO Word64)
-> (State# RealWorld -> (# State# RealWorld, Word64 #))
-> IO Word64
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#
i State# RealWorld
s of
      (# State# RealWorld
s', Int#
v #) -> (# State# RealWorld
s', Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int# -> Int
I# Int#
v) #)
{-# INLINE readBucketArray #-}


{- | Read all bucket values and atomically reset each to zero.
Uses fetch-and-add with the negated current value.
Between the read and the subtract another thread may increment,
but the net effect is correct across consecutive delta collections:
no sample is lost, at worst one sample shifts to the next cycle.

@since 0.0.1.0
-}
readAndResetBucketArray :: AtomicBucketArray -> IO (Vector Word64)
readAndResetBucketArray :: AtomicBucketArray -> IO (Vector Word64)
readAndResetBucketArray (AtomicBucketArray MutableByteArray# RealWorld
arr Int#
len) = do
  let !n :: Int
n = Int# -> Int
I# Int#
len
  Int -> (Int -> IO Word64) -> IO (Vector Word64)
forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> (Int -> m a) -> m (Vector a)
U.generateM Int
n ((Int -> IO Word64) -> IO (Vector Word64))
-> (Int -> IO Word64) -> IO (Vector Word64)
forall a b. (a -> b) -> a -> b
$ \(I# Int#
i) -> (State# RealWorld -> (# State# RealWorld, Word64 #)) -> IO Word64
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Word64 #)) -> IO Word64)
-> (State# RealWorld -> (# State# RealWorld, Word64 #))
-> IO Word64
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#
i State# RealWorld
s of
      (# State# RealWorld
s1, Int#
old #) -> 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#
i (Int# -> Int#
negateInt# Int#
old) State# RealWorld
s1 of
        (# State# RealWorld
s2, Int#
_ #) -> (# State# RealWorld
s2, Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int# -> Int
I# Int#
old) #)
{-# INLINE readAndResetBucketArray #-}