    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE MagicHash #-}
    -- |
    -- Module       : Data.ByteString.Base64.Internal
    -- Copyright    : (c) 2019-2023 Emily Pillmore
    -- License      : BSD-style
    --
    -- Maintainer   : Emily Pillmore <emilypi@cohomolo.gy>
    -- Stability    : Experimental
    -- Portability  : portable
    --
    -- Shared internal utils
    --
    module Data.ByteString.Base64.Internal.Utils
    ( EncodingTable(..)
    , aix
    , mask_2bits
    , mask_4bits
    , packTable
    , peekWord32BE
    , peekWord64BE
    , reChunkN
    , validateLastPos
    , w32
    , w64
    , w32_16
    , w64_16
    , writeNPlainForeignPtrBytes
    ) where

    import Data.Bits
    import Data.ByteString (ByteString)
    import qualified Data.ByteString as BS

    import Foreign.ForeignPtr
    import Foreign.Ptr
    import Foreign.Storable

    import GHC.ByteOrder
    import GHC.Exts
    import GHC.ForeignPtr
    import GHC.Word

    import System.IO.Unsafe

    -- | Only the lookup table need be a foreignptr,
    -- and then, only so that we can automate some touches to keep it alive
    --
    data EncodingTable = EncodingTable
      {-# UNPACK #-} !(Ptr Word8)
      {-# UNPACK #-} !(ForeignPtr Word16)

    -- | Read 'Word8' index off alphabet addr
    --
    aix :: Word8 -> Addr# -> Word8
    aix :: Word8 -> Addr# -> Word8
    aix Word8
    w8 Addr#
    alpha = Word8# -> Word8
    W8# (Addr# -> Int# -> Word8#
    indexWord8OffAddr# Addr#
    alpha Int#
    i)
      where
        !(I# Int#
    i) = Word8 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word8
    w8
    {-# INLINE aix #-}

    -- | Convert 'Word8''s into 'Word32''s
    --
    w32 :: Word8 -> Word32
    w32 :: Word8 -> Word32
    w32 = Word8 -> Word32
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral
    {-# INLINE w32 #-}

    -- | Convert 'Word8''s into 'Word32''s
    --
    w64 :: Word8 -> Word64
    w64 :: Word8 -> Word64
    w64 = Word8 -> Word64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral
    {-# INLINE w64 #-}

    -- | Convert 'Word8''s into 'Word32''s
    --
    w64_16 :: Word16 -> Word64
    w64_16 :: Word16 -> Word64
    w64_16 = Word16 -> Word64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral
    {-# INLINE w64_16 #-}

    w32_16 :: Word16 -> Word32
    w32_16 :: Word16 -> Word32
    w32_16 = Word16 -> Word32
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral
    {-# INLINE w32_16 #-}

    -- | Mask bottom 2 bits
    --
    mask_2bits :: Word8
    mask_2bits :: Word8
    mask_2bits = Word8
    3  -- (1 << 2) - 1
    {-# INLINE mask_2bits #-}

    -- | Mask bottom 4 bits
    --
    mask_4bits :: Word8
    mask_4bits :: Word8
    mask_4bits = Word8
    15 -- (1 << 4) - 1
    {-# INLINE mask_4bits #-}

    -- | Validate some ptr index against some bitmask
    --
    validateLastPos :: Word32 -> Word8 -> Bool
    validateLastPos :: Word32 -> Word8 -> Bool
    validateLastPos Word32
    pos Word8
    mask = (Word32 -> Word8
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word32
    pos Word8 -> Word8 -> Word8
    forall a. Bits a => a -> a -> a
    .&. Word8
    mask) Word8 -> Word8 -> Bool
    forall a. Eq a => a -> a -> Bool
    == Word8
    0
    {-# INLINE validateLastPos #-}

    -- | Allocate and fill @n@ bytes with some data
    --
    writeNPlainForeignPtrBytes
        :: ( Storable a
           , Storable b
           )
        => Int
        -> [a]
        -> ForeignPtr b
    writeNPlainForeignPtrBytes :: forall a b. (Storable a, Storable b) => Int -> [a] -> ForeignPtr b
    writeNPlainForeignPtrBytes !Int
    n [a]
    as = IO (ForeignPtr b) -> ForeignPtr b
    forall a. IO a -> a
    unsafeDupablePerformIO (IO (ForeignPtr b) -> ForeignPtr b)
    -> IO (ForeignPtr b) -> ForeignPtr b
    forall a b. (a -> b) -> a -> b
    $ do
        fp <- Int -> IO (ForeignPtr a)
    forall a. Int -> IO (ForeignPtr a)
    mallocPlainForeignPtrBytes Int
    n
        withForeignPtr fp $ \Ptr a
    p -> Ptr a -> [a] -> IO ()
    forall {b}. Storable b => Ptr b -> [b] -> IO ()
    go Ptr a
    p [a]
    as
        return (castForeignPtr fp)
      where
        go :: Ptr b -> [b] -> IO ()
    go !Ptr b
    _ [] = () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
        go !Ptr b
    p (b
    x:[b]
    xs) = Ptr b -> b -> IO ()
    forall a. Storable a => Ptr a -> a -> IO ()
    poke Ptr b
    p b
    x IO () -> IO () -> IO ()
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Ptr b -> [b] -> IO ()
    go (Ptr b -> Int -> Ptr b
    forall a b. Ptr a -> Int -> Ptr b
    plusPtr Ptr b
    p Int
    1) [b]
    xs

    -- | Pack an 'Addr#' into an encoding table of 'Word16's
    --
    packTable :: Addr# -> EncodingTable
    packTable :: Addr# -> EncodingTable
    packTable Addr#
    alphabet = EncodingTable
    etable
      where
        ix :: Int -> Word8
    ix (I# Int#
    n) = Word8# -> Word8
    W8# (Addr# -> Int# -> Word8#
    indexWord8OffAddr# Addr#
    alphabet Int#
    n)

        !etable :: EncodingTable
    etable =
          let bs :: [Word8]
    bs = [[Word8]] -> [Word8]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat
                [ [ Int -> Word8
    ix Int
    i, Int -> Word8
    ix Int
    j ]
                | !Int
    i <- [Int
    0..Int
    63]
                , !Int
    j <- [Int
    0..Int
    63]
                ]
          in Ptr Word8 -> ForeignPtr Word16 -> EncodingTable
    EncodingTable (Addr# -> Ptr Word8
    forall a. Addr# -> Ptr a
    Ptr Addr#
    alphabet) (Int -> [Word8] -> ForeignPtr Word16
    forall a b. (Storable a, Storable b) => Int -> [a] -> ForeignPtr b
    writeNPlainForeignPtrBytes Int
    8192 [Word8]
    bs)

    -- | Rechunk a list of bytestrings in multiples of 4
    --
    reChunkN :: Int -> [ByteString] -> [ByteString]
    reChunkN :: Int -> [ByteString] -> [ByteString]
    reChunkN Int
    n = [ByteString] -> [ByteString]
    go
      where
        go :: [ByteString] -> [ByteString]
    go [] = []
        go (ByteString
    b:[ByteString]
    bs) = case Int -> Int -> (Int, Int)
    forall a. Integral a => a -> a -> (a, a)
    divMod (ByteString -> Int
    BS.length ByteString
    b) Int
    n of
          (Int
    _, Int
    0) -> ByteString
    b ByteString -> [ByteString] -> [ByteString]
    forall a. a -> [a] -> [a]
    : [ByteString] -> [ByteString]
    go [ByteString]
    bs
          (Int
    d, Int
    _) -> case Int -> ByteString -> (ByteString, ByteString)
    BS.splitAt (Int
    d Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Int
    n) ByteString
    b of
            ~(ByteString
    h, ByteString
    t) -> ByteString
    h ByteString -> [ByteString] -> [ByteString]
    forall a. a -> [a] -> [a]
    : ByteString -> [ByteString] -> [ByteString]
    accum ByteString
    t [ByteString]
    bs

        accum :: ByteString -> [ByteString] -> [ByteString]
    accum ByteString
    acc [] = [ByteString
    acc]
        accum ByteString
    acc (ByteString
    c:[ByteString]
    cs) =
          case Int -> ByteString -> (ByteString, ByteString)
    BS.splitAt (Int
    n Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - ByteString -> Int
    BS.length ByteString
    acc) ByteString
    c of
            ~(ByteString
    h, ByteString
    t) ->
              let acc' :: ByteString
    acc' = ByteString -> ByteString -> ByteString
    BS.append ByteString
    acc ByteString
    h
              in if ByteString -> Int
    BS.length ByteString
    acc' Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    n
                 then
                   let cs' :: [ByteString]
    cs' = if ByteString -> Bool
    BS.null ByteString
    t then [ByteString]
    cs else ByteString
    t ByteString -> [ByteString] -> [ByteString]
    forall a. a -> [a] -> [a]
    : [ByteString]
    cs
                   in ByteString
    acc' ByteString -> [ByteString] -> [ByteString]
    forall a. a -> [a] -> [a]
    : [ByteString] -> [ByteString]
    go [ByteString]
    cs'
                 else ByteString -> [ByteString] -> [ByteString]
    accum ByteString
    acc' [ByteString]
    cs
    {-# INLINE reChunkN #-}

    peekWord32BE :: Ptr Word32 -> IO Word32
    peekWord32BE :: Ptr Word32 -> IO Word32
    peekWord32BE Ptr Word32
    p = case ByteOrder
    targetByteOrder of
      ByteOrder
    LittleEndian -> Word32 -> Word32
    byteSwap32 (Word32 -> Word32) -> IO Word32 -> IO Word32
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Ptr Word32 -> IO Word32
    forall a. Storable a => Ptr a -> IO a
    peek Ptr Word32
    p
      ByteOrder
    BigEndian -> Ptr Word32 -> IO Word32
    forall a. Storable a => Ptr a -> IO a
    peek Ptr Word32
    p
    {-# inline peekWord32BE #-}

    peekWord64BE :: Ptr Word64 -> IO Word64
    peekWord64BE :: Ptr Word64 -> IO Word64
    peekWord64BE Ptr Word64
    p = case ByteOrder
    targetByteOrder of
      ByteOrder
    LittleEndian -> Word64 -> Word64
    byteSwap64 (Word64 -> Word64) -> IO Word64 -> IO Word64
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Ptr Word64 -> IO Word64
    forall a. Storable a => Ptr a -> IO a
    peek Ptr Word64
    p
      ByteOrder
    BigEndian -> Ptr Word64 -> IO Word64
    forall a. Storable a => Ptr a -> IO a
    peek Ptr Word64
    p
    {-# inline peekWord64BE #-}