{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Data.ByteString.Base64.Arm
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- ARM NEON support for base64 encoding and decoding.

module Data.ByteString.Base64.Arm (
    base64_arm_available
  , encode
  , decode
  ) where

import qualified Data.Bits as B
import Data.Bits ((.&.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI
import Data.Word (Word8)
import Foreign.C.Types (CInt(..), CSize(..))
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peekElemOff)
import System.IO.Unsafe (unsafeDupablePerformIO)

-- ffi ------------------------------------------------------------------------

foreign import ccall unsafe "base64_encode_arm"
  c_base64_encode :: Ptr Word8 -> Ptr Word8 -> CSize -> IO ()

foreign import ccall unsafe "base64_decode_arm"
  c_base64_decode :: Ptr Word8 -> Ptr Word8 -> CSize -> CSize -> IO CInt

foreign import ccall unsafe "base64_arm_available"
  c_base64_arm_available :: IO CInt

-- utilities ------------------------------------------------------------------

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

-- api ------------------------------------------------------------------------

-- | Are ARM NEON extensions available?
base64_arm_available :: Bool
base64_arm_available :: Bool
base64_arm_available =
  IO CInt -> CInt
forall a. IO a -> a
unsafeDupablePerformIO IO CInt
c_base64_arm_available CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
/= CInt
0
{-# NOINLINE base64_arm_available #-}

-- | Encode a base256 'ByteString' as base64 using NEON.
encode :: BS.ByteString -> BS.ByteString
encode :: ByteString -> ByteString
encode (BI.PS ForeignPtr Word8
sfp Int
soff Int
l) =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
BI.unsafeCreate ((Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4) ((Ptr Word8 -> IO ()) -> ByteString)
-> (Ptr Word8 -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dst ->
    ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp0 ->
      Ptr Word8 -> Ptr Word8 -> CSize -> IO ()
c_base64_encode (Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff) Ptr Word8
dst (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fi Int
l)

-- | Decode a base64 'ByteString' to base256 using NEON.  Returns
--   'Nothing' on malformed input.
decode :: BS.ByteString -> Maybe BS.ByteString
decode :: ByteString -> Maybe ByteString
decode (BI.PS ForeignPtr Word8
sfp Int
soff Int
l)
  | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0          = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
BS.empty
  | Int
l Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x03 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Maybe ByteString
forall a. Maybe a
Nothing
  | Bool
otherwise = IO (Maybe ByteString) -> Maybe ByteString
forall a. IO a -> a
unsafeDupablePerformIO (IO (Maybe ByteString) -> Maybe ByteString)
-> IO (Maybe ByteString) -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$
      ForeignPtr Word8
-> (Ptr Word8 -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO (Maybe ByteString)) -> IO (Maybe ByteString))
-> (Ptr Word8 -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp0 -> do
        let !sp :: Ptr Word8
sp = Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff :: Ptr Word8
        c_pre <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)
        c_end <- peekElemOff sp (l - 1)
        let !pad_pre = Word8
c_pre Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x3D
            !pad_end = Word8
c_end Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x3D
        if pad_pre && not pad_end
          then pure Nothing
          else do
            let !pad = (if Bool
pad_pre then Int
2 else if Bool
pad_end then Int
1 else Int
0)
                     :: Int
                !nfull  = Int
l Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
2
                !outlen = Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
pad
            fp <- BI.mallocByteString outlen
            ok <- withForeignPtr fp $ \Ptr Word8
dst ->
              Ptr Word8 -> Ptr Word8 -> CSize -> CSize -> IO CInt
c_base64_decode Ptr Word8
sp Ptr Word8
dst (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fi Int
l) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fi Int
outlen)
            pure $! if ok /= 0
              then Just (BI.PS fp 0 outlen)
              else Nothing