{-# LANGUAGE BangPatterns #-}

module DataFrame.Internal.Binary where

import Data.Bits (Bits (unsafeShiftL, (.|.)))
import Data.ByteString (toStrict)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (toLazyByteString, word32LE, word64LE)
import qualified Data.ByteString.Unsafe as BS
import Data.Int (Int32)
import Data.Word (Word32, Word64, Word8)

littleEndianWord32 :: BS.ByteString -> Word32
littleEndianWord32 :: ByteString -> Word32
littleEndianWord32 ByteString
bytes
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
4 =
        Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleWord32
            (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bytes Int
0)
            (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bytes Int
1)
            (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bytes Int
2)
            (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bytes Int
3)
    | Bool
otherwise =
        Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleWord32
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
0)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
1)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
2)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
3)
  where
    len :: Int
len = ByteString -> Int
BS.length ByteString
bytes
{-# INLINE littleEndianWord32 #-}

littleEndianWord64 :: BS.ByteString -> Word64
littleEndianWord64 :: ByteString -> Word64
littleEndianWord64 ByteString
bytes
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
8 =
        Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word64
assembleWord64
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
0)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
1)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
2)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
3)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
4)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
5)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
6)
            (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bytes Int
7)
    | Bool
otherwise =
        Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word64
assembleWord64
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
0)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
1)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
2)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
3)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
4)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
5)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
6)
            (Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
7)
  where
    len :: Int
len = ByteString -> Int
BS.length ByteString
bytes
{-# INLINE littleEndianWord64 #-}

littleEndianInt32 :: BS.ByteString -> Int32
littleEndianInt32 :: ByteString -> Int32
littleEndianInt32 = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Int32) -> (ByteString -> Word32) -> ByteString -> Int32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Word32
littleEndianWord32
{-# INLINE littleEndianInt32 #-}

word64ToLittleEndian :: Word64 -> BS.ByteString
word64ToLittleEndian :: Word64 -> ByteString
word64ToLittleEndian = LazyByteString -> ByteString
toStrict (LazyByteString -> ByteString)
-> (Word64 -> LazyByteString) -> Word64 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
toLazyByteString (Builder -> LazyByteString)
-> (Word64 -> Builder) -> Word64 -> LazyByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Builder
word64LE
{-# INLINE word64ToLittleEndian #-}

word32ToLittleEndian :: Word32 -> BS.ByteString
word32ToLittleEndian :: Word32 -> ByteString
word32ToLittleEndian = LazyByteString -> ByteString
toStrict (LazyByteString -> ByteString)
-> (Word32 -> LazyByteString) -> Word32 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
toLazyByteString (Builder -> LazyByteString)
-> (Word32 -> Builder) -> Word32 -> LazyByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Builder
word32LE
{-# INLINE word32ToLittleEndian #-}

byteAtOrZero :: Int -> BS.ByteString -> Int -> Word8
byteAtOrZero :: Int -> ByteString -> Int -> Word8
byteAtOrZero Int
len ByteString
bytes Int
i
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bytes Int
i
    | Bool
otherwise = Word8
0
{-# INLINE byteAtOrZero #-}

assembleWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleWord32 !Word8
b0 !Word8
b1 !Word8
b2 !Word8
b3 =
    Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b0
        Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b1 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8)
        Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b2 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
16)
        Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b3 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
24)
{-# INLINE assembleWord32 #-}

assembleWord64 ::
    Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word64
assembleWord64 :: Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word8
-> Word64
assembleWord64 !Word8
b0 !Word8
b1 !Word8
b2 !Word8
b3 !Word8
b4 !Word8
b5 !Word8
b6 !Word8
b7 =
    Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b0
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b2 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
16)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b3 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
24)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b4 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
32)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b5 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
40)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b6 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
48)
        Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b7 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
56)
{-# INLINE assembleWord64 #-}