{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}
-- |
-- Module       : Data.ByteString.Base64.Internal.Tail
-- Copyright    : (c) 2019-2023 Emily Pillmore
-- License      : BSD-style
--
-- Maintainer   : Emily Pillmore <emilypi@cohomolo.gy>
-- Stability    : Experimental
-- Portability  : portable
--
-- Finalizers for the encoding loop
--
module Data.ByteString.Base64.Internal.Tail
( loopTail
, loopTailNoPad
) where

import Data.Bits
import Data.ByteString.Internal
import Data.ByteString.Base64.Internal.Utils

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

import GHC.Exts
import GHC.Word

-- | Finalize an encoded bytestring by filling in the remaining
-- bytes and any padding
--
loopTail
    :: ForeignPtr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> IO ByteString
loopTail :: ForeignPtr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> IO ByteString
loopTail !ForeignPtr Word8
dfp !Ptr Word8
dptr (Ptr !Addr#
alpha) !Ptr Word8
end !Ptr Word8
sptr_ !Ptr Word8
dptr_ = Ptr Word8 -> Ptr Word8 -> IO ByteString
go Ptr Word8
sptr_ Ptr Word8
dptr_
  where
    go :: Ptr Word8 -> Ptr Word8 -> IO ByteString
go Ptr Word8
src Ptr Word8
dst
      | Ptr Word8
src Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
end =
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
PS ForeignPtr Word8
dfp Int
0 (Ptr Word8
dst Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
dptr)
      | Bool
otherwise = do
        !x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src
        let !a = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR Word8
x Int
2
            !carry0 = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03) Int
4

        -- poke first 6 bits
        poke @Word8 dst $ aix a alpha

        if src `plusPtr` 1 == end then do
          -- if no other bytes, poke carry bits
          poke @Word8 (dst `plusPtr` 1) $ aix carry0 alpha
          poke @Word8 (dst `plusPtr` 2) 0x3d
          poke @Word8 (dst `plusPtr` 3) 0x3d

        else do
          !y <- peek @Word8 $ src `plusPtr` 1

          let !b = Word8
carry0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word8
y Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xf0) Int
4
              !carry1 = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8
y Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0f) Int
2

          poke @Word8 (dst `plusPtr` 1) $ aix b alpha
          poke @Word8 (dst `plusPtr` 2) $ aix carry1 alpha
          poke @Word8 (dst `plusPtr` 3) 0x3d

        pure $ PS dfp 0 (4 + minusPtr dst dptr)
{-# inline loopTail #-}

-- | Finalize a bytestring by filling out the remaining bits
-- without padding.
--
loopTailNoPad
    :: ForeignPtr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> Ptr Word8
    -> IO ByteString
loopTailNoPad :: ForeignPtr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> IO ByteString
loopTailNoPad !ForeignPtr Word8
dfp (Ptr !Addr#
alpha) !Ptr Word8
dptr !Ptr Word8
end !Ptr Word8
src !Ptr Word8
dst
    | Ptr Word8
src Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
end = ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr Word8 -> Int -> Int -> ByteString
PS ForeignPtr Word8
dfp Int
0 (Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
minusPtr Ptr Word8
dst Ptr Word8
dptr))
    | Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
src Int
1 Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
end = do
      !x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src

      let !a = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xfc) Int
2
          !b = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03) Int
4

      poke @Word8 dst (aix a alpha)
      poke @Word8 (plusPtr dst 1) (aix b alpha)
      return (PS dfp 0 (2 + minusPtr dst dptr))
    | Bool
otherwise = do
      !x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src
      !y <- peek @Word8 (plusPtr src 1)

      let !a = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xfc) Int
2
          !b = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03) Int
4

      let !c = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word8
y Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xf0) Int
4 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
b
          !d = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8
y Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0f) Int
2

      poke @Word8 dst (aix a alpha)
      poke @Word8 (plusPtr dst 1) (aix c alpha)
      poke @Word8 (plusPtr dst 2) (aix d alpha)
      return (PS dfp 0 (3 + minusPtr dst dptr))
{-# inline loopTailNoPad #-}