{-# LANGUAGE BangPatterns #-}
-- |
-- Module       : Data.ByteString.Base64.Internal.W64.Loop
-- Copyright    : (c) 2019-2023 Emily Pillmore
-- License      : BSD-style
--
-- Maintainer   : Emily Pillmore <emilypi@cohomolo.gy>
-- Stability    : Experimental
-- Portability  : portable
--
-- 'Word64'-optimized inner loops
--
module Data.ByteString.Base64.Internal.W64.Loop
( innerLoop
, decodeLoop
, decodeLoopNoError
, lenientLoop
) where


import Data.Bits
import Data.ByteString.Internal
import Data.ByteString.Base64.Internal.Utils
import qualified Data.ByteString.Base64.Internal.W16.Loop as W16
import Data.Text (Text)

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

import GHC.Word


-- | Encoding inner loop. Packs 6 bytes from src pointer into
-- the first 6 bits of 4 'Word12''s (using the encoding table,
-- as 2 'Word12''s ), writing these to the dst pointer.
--
innerLoop
    :: Ptr Word16
    -> Ptr Word64
    -> Ptr Word64
    -> Ptr Word64
    -> (Ptr Word8 -> Ptr Word8 -> IO ByteString)
    -> IO ByteString
innerLoop :: Ptr Word16
-> Ptr Word64
-> Ptr Word64
-> Ptr Word64
-> (Ptr Word8 -> Ptr Word8 -> IO ByteString)
-> IO ByteString
innerLoop !Ptr Word16
etable !Ptr Word64
sptr !Ptr Word64
dptr !Ptr Word64
end Ptr Word8 -> Ptr Word8 -> IO ByteString
finish = Ptr Word64 -> Ptr Word64 -> IO ByteString
forall {b}.
(Storable b, Num b) =>
Ptr Word64 -> Ptr b -> IO ByteString
go Ptr Word64
sptr Ptr Word64
dptr
  where
    go :: Ptr Word64 -> Ptr b -> IO ByteString
go !Ptr Word64
src !Ptr b
dst
      | Ptr Word64 -> Int -> Ptr Word64
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word64
src Int
7 Ptr Word64 -> Ptr Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr Word64
end =
        Ptr Word16
-> Ptr Word8
-> Ptr Word16
-> Ptr Word8
-> (Ptr Word8 -> Ptr Word8 -> IO ByteString)
-> IO ByteString
W16.innerLoop Ptr Word16
etable (Ptr Word64 -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr Word64
src) (Ptr b -> Ptr Word16
forall a b. Ptr a -> Ptr b
castPtr Ptr b
dst) (Ptr Word64 -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr Word64
end) Ptr Word8 -> Ptr Word8 -> IO ByteString
finish
      | Bool
otherwise = do
        !t <- Ptr Word64 -> IO Word64
peekWord64BE Ptr Word64
src

        let !a = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
t Int
52 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xfff
            !b = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
t Int
40 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xfff
            !c = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
t Int
28 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xfff
            !d = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
t Int
16 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xfff

        !w <- w64_16 <$> peekElemOff etable (fromIntegral a)
        !x <- w64_16 <$> peekElemOff etable (fromIntegral b)
        !y <- w64_16 <$> peekElemOff etable (fromIntegral c)
        !z <- w64_16 <$> peekElemOff etable (fromIntegral d)

        let !xx = Word64
w
               Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftL Word64
x Int
16
               Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftL Word64
y Int
32
               Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftL Word64
z Int
48

        poke dst (fromIntegral xx)

        go (plusPtr src 6) (plusPtr dst 8)
{-# inline innerLoop #-}

decodeLoop
    :: Ptr Word8
        -- ^ decode lookup table
    -> Ptr Word8
        -- ^ src pointer
    -> Ptr Word8
        -- ^ dst pointer
    -> Ptr Word8
        -- ^ end of src ptr
    -> ForeignPtr Word8
        -- ^ dst foreign ptr (for consing bs)
    -> IO (Either Text ByteString)
decodeLoop :: Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO (Either Text ByteString)
decodeLoop = Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO (Either Text ByteString)
W16.decodeLoop
{-# inline decodeLoop #-}

decodeLoopNoError
    :: Ptr Word8
        -- ^ decode lookup table
    -> Ptr Word8
        -- ^ src pointer
    -> Ptr Word8
        -- ^ dst pointer
    -> Ptr Word8
        -- ^ end of src ptr
    -> ForeignPtr Word8
        -- ^ dst foreign ptr (for consing bs)
    -> IO ByteString
decodeLoopNoError :: Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO ByteString
decodeLoopNoError = Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO ByteString
W16.decodeLoopNoError
{-# inline decodeLoopNoError #-}

lenientLoop
    :: Ptr Word8
        -- ^ decode lookup table
    -> Ptr Word8
        -- ^ src pointer
    -> Ptr Word8
        -- ^ dst pointer
    -> Ptr Word8
        -- ^ end of src ptr
    -> ForeignPtr Word8
        -- ^ dst foreign ptr (for consing bs)
    -> IO ByteString
lenientLoop :: Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO ByteString
lenientLoop = Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> Ptr Word8
-> ForeignPtr Word8
-> IO ByteString
W16.lenientLoop