{-# 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
        !Word8
x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src
        let !a :: Word8
a = Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR Word8
x Int
2
            !carry0 :: Word8
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
        forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 Ptr Word8
dst (Word8 -> IO ()) -> Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Word8 -> Addr# -> Word8
aix Word8
a Addr#
alpha

        if Ptr Word8
src Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1 Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
end then do
          -- if no other bytes, poke carry bits
          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1) (Word8 -> IO ()) -> Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Word8 -> Addr# -> Word8
aix Word8
carry0 Addr#
alpha
          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
2) Word8
0x3d
          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
3) Word8
0x3d

        else do
          !Word8
y <- forall a. Storable a => Ptr a -> IO a
peek @Word8 (Ptr Word8 -> IO Word8) -> Ptr Word8 -> IO Word8
forall a b. (a -> b) -> a -> b
$ Ptr Word8
src Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1

          let !b :: Word8
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
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

          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1) (Word8 -> IO ()) -> Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Word8 -> Addr# -> Word8
aix Word8
b Addr#
alpha
          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
2) (Word8 -> IO ()) -> Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Word8 -> Addr# -> Word8
aix Word8
carry1 Addr#
alpha
          forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
3) Word8
0x3d

        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 (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
minusPtr Ptr Word8
dst Ptr Word8
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
      !Word8
x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src

      let !a :: Word8
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
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

      forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 Ptr Word8
dst (Word8 -> Addr# -> Word8
aix Word8
a Addr#
alpha)
      forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
dst Int
1) (Word8 -> Addr# -> Word8
aix Word8
b Addr#
alpha)
      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 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
minusPtr Ptr Word8
dst Ptr Word8
dptr))
    | Bool
otherwise = do
      !Word8
x <- forall a. Storable a => Ptr a -> IO a
peek @Word8 Ptr Word8
src
      !Word8
y <- forall a. Storable a => Ptr a -> IO a
peek @Word8 (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
src Int
1)

      let !a :: Word8
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
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
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
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

      forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 Ptr Word8
dst (Word8 -> Addr# -> Word8
aix Word8
a Addr#
alpha)
      forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
dst Int
1) (Word8 -> Addr# -> Word8
aix Word8
c Addr#
alpha)
      forall a. Storable a => Ptr a -> a -> IO ()
poke @Word8 (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
dst Int
2) (Word8 -> Addr# -> Word8
aix Word8
d Addr#
alpha)
      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 (Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
minusPtr Ptr Word8
dst Ptr Word8
dptr))
{-# inline loopTailNoPad #-}