{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}

{- | Conversion of the hs-opentelemetry internal representation of the trace ID and the span ID and the Datadog header representation of them each other.

+----------+-----------------+----------------+
|          | Trace ID        | Span ID        |
+----------+-----------------+----------------+
| Internal | 128-bit integer | 64-bit integer |
+----------+-----------------+----------------+
| Datadog  | ASCII text of   | ASCII text of  |
| Header   | 64-bit integer  | 64-bit integer |
+----------+-----------------+----------------+
-}
module OpenTelemetry.Propagator.Datadog.Internal (
  newTraceIdFromHeader,
  newSpanIdFromHeader,
  newHeaderFromTraceId,
  newHeaderFromSpanId,
  indexByteArrayNbo,
) where

import Data.Bits (Bits (shift))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as BL
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short as SB
import qualified Data.ByteString.Short.Internal as SBI
import qualified Data.Char as C
import Data.Primitive.ByteArray (ByteArray (ByteArray), indexByteArray)
import Data.Primitive.Ptr (writeOffPtr)
import Data.Word (Word64, Word8)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Storable (peekElemOff)
import System.IO.Unsafe (unsafeDupablePerformIO)


newTraceIdFromHeader
  :: ByteString
  -- ^ ASCII text of 64-bit integer
  -> ShortByteString
  -- ^ 128-bit integer
newTraceIdFromHeader :: ByteString -> ShortByteString
newTraceIdFromHeader ByteString
bs =
  let w64 :: Word64
w64 = ByteString -> Word64
readWord64BS ByteString
bs
      builder :: Builder
builder = Word64 -> Builder
BB.word64BE Word64
0 forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
BB.word64BE Word64
w64
   in ByteString -> ShortByteString
SB.toShort forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
BB.toLazyByteString Builder
builder


newSpanIdFromHeader
  :: ByteString
  -- ^ ASCII text of 64-bit integer
  -> ShortByteString
  -- ^ 64-bit integer
newSpanIdFromHeader :: ByteString -> ShortByteString
newSpanIdFromHeader ByteString
bs =
  let w64 :: Word64
w64 = ByteString -> Word64
readWord64BS ByteString
bs
      builder :: Builder
builder = Word64 -> Builder
BB.word64BE Word64
w64
   in ByteString -> ShortByteString
SB.toShort forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
BB.toLazyByteString Builder
builder


readWord64BS :: ByteString -> Word64
readWord64BS :: ByteString -> Word64
readWord64BS (BI.PS ForeignPtr Word8
fptr Int
_ Int
len) =
  -- Safe.
  forall a. IO a -> a
unsafeDupablePerformIO forall a b. (a -> b) -> a -> b
$
    forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr Ptr Word8 -> IO Word64
readWord64Ptr
  where
    readWord64Ptr :: Ptr Word8 -> IO Word64
readWord64Ptr Ptr Word8
ptr =
      Int -> Word64 -> IO Word64
readWord64PtrOffset Int
0 Word64
0
      where
        readWord64PtrOffset :: Int -> Word64 -> IO Word64
readWord64PtrOffset Int
offset Word64
acc
          | Int
offset forall a. Ord a => a -> a -> Bool
< Int
len = do
              Word8
b <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
ptr Int
offset
              let n :: Word64
n = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Word8 -> Word8
asciiWord8ToWord8 Word8
b :: Word64
              Int -> Word64 -> IO Word64
readWord64PtrOffset (Int
offset forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ Word64
n forall a. Num a => a -> a -> a
+ Word64
acc forall a. Num a => a -> a -> a
* Word64
10
          | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Word64
acc


asciiWord8ToWord8 :: Word8 -> Word8
asciiWord8ToWord8 :: Word8 -> Word8
asciiWord8ToWord8 Word8
b = Word8
b forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
C.ord Char
'0')


newHeaderFromTraceId
  :: ShortByteString
  -- ^ 128-bit integer
  -> ByteString
  -- ^ ASCII text of 64-bit integer
newHeaderFromTraceId :: ShortByteString -> ByteString
newHeaderFromTraceId (SBI.SBS ByteArray#
ba) =
  let w64 :: Word64
w64 = ByteArray -> Int -> Word64
indexByteArrayNbo (ByteArray# -> ByteArray
ByteArray ByteArray#
ba) Int
1
   in Word64 -> ByteString
showWord64BS Word64
w64


newHeaderFromSpanId
  :: ShortByteString
  -- ^ 64-bit integer
  -> ByteString
  -- ^ ASCII text of 64-bit integer
newHeaderFromSpanId :: ShortByteString -> ByteString
newHeaderFromSpanId (SBI.SBS ByteArray#
ba) =
  let w64 :: Word64
w64 = ByteArray -> Int -> Word64
indexByteArrayNbo (ByteArray# -> ByteArray
ByteArray ByteArray#
ba) Int
0
   in Word64 -> ByteString
showWord64BS Word64
w64


-- | Read 'ByteArray' to 'Word64' with network-byte-order.
indexByteArrayNbo
  :: ByteArray
  -> Int
  -- ^ Offset in 'Word64'-size unit
  -> Word64
indexByteArrayNbo :: ByteArray -> Int -> Word64
indexByteArrayNbo ByteArray
ba Int
offset =
  Int -> Word64 -> Word64
loop Int
0 Word64
0
  where
    loop :: Int -> Word64 -> Word64
loop Int
8 Word64
acc = Word64
acc
    loop Int
n Word64
acc = Int -> Word64 -> Word64
loop (Int
n forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ forall a. Bits a => a -> Int -> a
shift Word64
acc Int
8 forall a. Num a => a -> a -> a
+ Word8 -> Word64
word8ToWord64 (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
ba forall a b. (a -> b) -> a -> b
$ Int
8 forall a. Num a => a -> a -> a
* Int
offset forall a. Num a => a -> a -> a
+ Int
n)


showWord64BS :: Word64 -> ByteString
showWord64BS :: Word64 -> ByteString
showWord64BS Word64
v =
  -- Safe.
  forall a. IO a -> a
unsafeDupablePerformIO forall a b. (a -> b) -> a -> b
$
    Int -> (Ptr Word8 -> IO Int) -> IO ByteString
BI.createUptoN Int
20 forall {m :: * -> *}. PrimMonad m => Ptr Word8 -> m Int
writeWord64Ptr -- 20 = length (show (maxBound :: Word64))
  where
    writeWord64Ptr :: Ptr Word8 -> m Int
writeWord64Ptr Ptr Word8
ptr =
      forall {m :: * -> *} {a} {t}.
(PrimMonad m, Integral a, Integral t) =>
t -> a -> Int -> Bool -> m Int
loop (Int
19 :: Int) Word64
v Int
0 Bool
False
      where
        loop :: t -> a -> Int -> Bool -> m Int
loop t
0 a
v Int
offset Bool
_ = do
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Ptr a -> Int -> a -> m ()
writeOffPtr Ptr Word8
ptr Int
offset (Word8 -> Word8
word8ToAsciiWord8 forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
v)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
offset forall a. Num a => a -> a -> a
+ Int
1
        loop t
n a
v Int
offset Bool
upper = do
          let (a
p, a
q) = a
v forall a. Integral a => a -> a -> (a, a)
`divMod` (a
10 forall a b. (Num a, Integral b) => a -> b -> a
^ t
n)
          if a
p forall a. Eq a => a -> a -> Bool
== a
0 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
upper
            then t -> a -> Int -> Bool -> m Int
loop (t
n forall a. Num a => a -> a -> a
- t
1) a
q Int
offset Bool
upper
            else do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Ptr a -> Int -> a -> m ()
writeOffPtr Ptr Word8
ptr Int
offset (Word8 -> Word8
word8ToAsciiWord8 forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
p)
              t -> a -> Int -> Bool -> m Int
loop (t
n forall a. Num a => a -> a -> a
- t
1) a
q (Int
offset forall a. Num a => a -> a -> a
+ Int
1) Bool
True


word8ToAsciiWord8 :: Word8 -> Word8
word8ToAsciiWord8 :: Word8 -> Word8
word8ToAsciiWord8 Word8
b = Word8
b forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
C.ord Char
'0')


word8ToWord64 :: Word8 -> Word64
word8ToWord64 :: Word8 -> Word64
word8ToWord64 = forall a b. (Integral a, Num b) => a -> b
fromIntegral