{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
-- |
-- Module       : Data.ByteString.Base64.Internal
-- Copyright    : (c) 2019-2020 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 (W8# Word#
i) Addr#
alpha = Word# -> Word8
W8# (Addr# -> Int# -> Word#
indexWord8OffAddr# Addr#
alpha (Word# -> Int#
word2Int# Word#
i))
{-# 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 :: 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
    ForeignPtr a
fp <- Int -> IO (ForeignPtr a)
forall a. Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytes Int
n
    ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fp ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> Ptr a -> [a] -> IO ()
forall b. Storable b => Ptr b -> [b] -> IO ()
go Ptr a
p [a]
as
    ForeignPtr b -> IO (ForeignPtr b)
forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr a -> ForeignPtr b
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr a
fp)
  where
    go :: Ptr b -> [b] -> IO ()
go !Ptr b
_ [] = () -> IO ()
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 (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) = Word# -> Word8
W8# (Addr# -> Int# -> Word#
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 #-}