{-# 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 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
BB.word64BE Word64
w64
  in ByteString -> ShortByteString
SB.toShort (ByteString -> ShortByteString) -> ByteString -> ShortByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
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 (ByteString -> ShortByteString) -> ByteString -> ShortByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
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.
  IO Word64 -> Word64
forall a. IO a -> a
unsafeDupablePerformIO (IO Word64 -> Word64) -> IO Word64 -> Word64
forall a b. (a -> b) -> a -> b
$
    ForeignPtr Word8 -> (Ptr Word8 -> IO Word64) -> IO Word64
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do
              Word8
b <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
ptr Int
offset
              let n :: Word64
n = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word64) -> Word8 -> Word64
forall a b. (a -> b) -> a -> b
$ Word8 -> Word8
asciiWord8ToWord8 Word8
b :: Word64
              Int -> Word64 -> IO Word64
readWord64PtrOffset (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word64 -> IO Word64) -> Word64 -> IO Word64
forall a b. (a -> b) -> a -> b
$ Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
acc Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
10
          | Bool
otherwise = Word64 -> IO Word64
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word64
acc


asciiWord8ToWord8 :: Word8 -> Word8
asciiWord8ToWord8 :: Word8 -> Word8
asciiWord8ToWord8 Word8
b = Word8
b Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Int -> Word8
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word64 -> Word64) -> Word64 -> Word64
forall a b. (a -> b) -> a -> b
$ Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
shift Word64
acc Int
8 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word8 -> Word64
word8ToWord64 (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
ba (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)


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


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


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