{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}

module Data.Bytes.Encode.LittleEndian
  ( word16
  , word32
  , word64
  , int16
  , int32
  , int64
  ) where

import Control.Monad.ST.Run (runByteArrayST)
import Data.Bits (unsafeShiftR)
import Data.Bytes.Types (Bytes)
import Data.Int (Int16, Int32, Int64)
import Data.Primitive (ByteArray)
import Data.Word (Word16, Word32, Word64, Word8)

import qualified Data.Bytes.Pure as Pure
import qualified Data.Primitive as PM

-- | Encode a 32-bit signed integer as 4 bytes.
int32 :: Int32 -> Bytes
{-# INLINE int32 #-}
int32 :: Int32 -> Bytes
int32 = Word32 -> Bytes
word32 (Word32 -> Bytes) -> (Int32 -> Word32) -> Int32 -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int32 @Word32

-- | Encode a 32-bit unsigned integer as 4 bytes.
word32 :: Word32 -> Bytes
word32 :: Word32 -> Bytes
word32 !Word32
w = ByteArray -> Bytes
Pure.fromByteArray (Word32 -> ByteArray
word32U Word32
w)

word32U :: Word32 -> ByteArray
word32U :: Word32 -> ByteArray
word32U !Word32
w = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  arr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
4
  PM.writeByteArray arr 3 (fromIntegral @Word32 @Word8 (unsafeShiftR w 24))
  PM.writeByteArray arr 2 (fromIntegral @Word32 @Word8 (unsafeShiftR w 16))
  PM.writeByteArray arr 1 (fromIntegral @Word32 @Word8 (unsafeShiftR w 8))
  PM.writeByteArray arr 0 (fromIntegral @Word32 @Word8 w)
  PM.unsafeFreezeByteArray arr

-- | Encode a 16-bit signed integer as 4 bytes.
int16 :: Int16 -> Bytes
{-# INLINE int16 #-}
int16 :: Int16 -> Bytes
int16 = Word16 -> Bytes
word16 (Word16 -> Bytes) -> (Int16 -> Word16) -> Int16 -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int16 @Word16

-- | Encode a 16-bit unsigned integer as 4 bytes.
word16 :: Word16 -> Bytes
word16 :: Word16 -> Bytes
word16 !Word16
w = ByteArray -> Bytes
Pure.fromByteArray (Word16 -> ByteArray
word16U Word16
w)

word16U :: Word16 -> ByteArray
word16U :: Word16 -> ByteArray
word16U !Word16
w = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  arr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
2
  PM.writeByteArray arr 1 (fromIntegral @Word16 @Word8 (unsafeShiftR w 8))
  PM.writeByteArray arr 0 (fromIntegral @Word16 @Word8 w)
  PM.unsafeFreezeByteArray arr

-- | Encode a 16-bit signed integer as 4 bytes.
int64 :: Int64 -> Bytes
{-# INLINE int64 #-}
int64 :: Int64 -> Bytes
int64 = Word64 -> Bytes
word64 (Word64 -> Bytes) -> (Int64 -> Word64) -> Int64 -> Bytes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int64 @Word64

-- | Encode a 16-bit unsigned integer as 4 bytes.
word64 :: Word64 -> Bytes
word64 :: Word64 -> Bytes
word64 !Word64
w = ByteArray -> Bytes
Pure.fromByteArray (Word64 -> ByteArray
word64U Word64
w)

word64U :: Word64 -> ByteArray
word64U :: Word64 -> ByteArray
word64U !Word64
w = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  arr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
8
  PM.writeByteArray arr 7 (fromIntegral @Word64 @Word8 (unsafeShiftR w 56))
  PM.writeByteArray arr 6 (fromIntegral @Word64 @Word8 (unsafeShiftR w 48))
  PM.writeByteArray arr 5 (fromIntegral @Word64 @Word8 (unsafeShiftR w 40))
  PM.writeByteArray arr 4 (fromIntegral @Word64 @Word8 (unsafeShiftR w 32))
  PM.writeByteArray arr 3 (fromIntegral @Word64 @Word8 (unsafeShiftR w 24))
  PM.writeByteArray arr 2 (fromIntegral @Word64 @Word8 (unsafeShiftR w 16))
  PM.writeByteArray arr 1 (fromIntegral @Word64 @Word8 (unsafeShiftR w 8))
  PM.writeByteArray arr 0 (fromIntegral @Word64 @Word8 w)
  PM.unsafeFreezeByteArray arr