{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module: Data.ByteString.Base64
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Pure base64 encoding and decoding of strict bytestrings.

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

import qualified Data.Bits as B
import Data.Bits ((.&.), (.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64.Arm as Arm
import qualified Data.ByteString.Internal as BI
import Data.Word (Word8)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (peekElemOff, pokeElemOff)
import System.IO.Unsafe (unsafeDupablePerformIO)

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

-- 64-byte table.  Indexed by 6-bit value (0..63), yields the
-- corresponding base64 alphabet character.  All-ASCII content means
-- the bytestring 'IsString' rule rewrites this to 'unsafePackAddress'
-- and the bytes live in static rodata.
enc_tab :: BS.ByteString
enc_tab :: ByteString
enc_tab =
  ByteString
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
{-# NOINLINE enc_tab #-}

-- 256-byte table.  Index by an ASCII byte to obtain its 6-bit value;
-- valid base64 chars ('A'..'Z', 'a'..'z', '0'..'9', '+', '/') map to
-- 0x40..0x7f, every other byte (including '=') maps to 0x80.
--
-- The encoding is chosen so the literal is strictly ASCII and contains
-- no embedded NUL, which is what the bytestring 'IsString' rule needs
-- to rewrite it into 'unsafePackAddress' (cf. 'enc_tab') — the bytes
-- end up in static rodata, with no CAF allocation.
--
-- The 0x80 sentinel is distinguished by bit 7; no value 0x40..0x7f
-- carries that bit, so 'decode' OR-folds every lookup into an
-- accumulator and tests 'acc .&. 0x80 == 0' once at the end.  The
-- low 6 bits of each entry are the 6-bit value, possibly contaminated
-- by the 0x40 flag bit; the b0/b1/b2 formulas mask each subexpression
-- before combining so the flag never bleeds into the output bytes.
dec_tab :: BS.ByteString
dec_tab :: ByteString
dec_tab =
  ByteString
"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x7E\x80\x80\x80\x7F\
  \\x74\x75\x76\x77\x78\x79\x7A\x7B\x7C\x7D\x80\x80\x80\x80\x80\x80\
  \\x80\x40\x41\x42\x43\x44\x45\x46\x47\x48\x49\x4A\x4B\x4C\x4D\x4E\
  \\x4F\x50\x51\x52\x53\x54\x55\x56\x57\x58\x59\x80\x80\x80\x80\x80\
  \\x80\x5A\x5B\x5C\x5D\x5E\x5F\x60\x61\x62\x63\x64\x65\x66\x67\x68\
  \\x69\x6A\x6B\x6C\x6D\x6E\x6F\x70\x71\x72\x73\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\
  \\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80"
{-# NOINLINE dec_tab #-}

-- | Encode a base256 'ByteString' as base64.
--
--   Uses ARM NEON extensions when available, otherwise a pure
--   Haskell scalar loop.
--
--   >>> encode "hello world"
--   "aGVsbG8gd29ybGQ="
encode :: BS.ByteString -> BS.ByteString
encode :: ByteString -> ByteString
encode ByteString
bs
  | Bool
Arm.base64_arm_available = ByteString -> ByteString
Arm.encode ByteString
bs
  | Bool
otherwise                = ByteString -> ByteString
encode_scalar ByteString
bs
{-# INLINABLE encode #-}

-- | Decode a base64 'ByteString' to base256.
--
--   Uses ARM NEON extensions when available, otherwise a pure
--   Haskell scalar loop.  Invalid inputs (including incorrectly-
--   padded or non-canonical inputs) will produce 'Nothing'.
--
--   >>> decode "aGVsbG8gd29ybGQ="
--   Just "hello world"
--   >>> decode "aGVsbG8gd29ybGQ" -- missing padding
--   Nothing
decode :: BS.ByteString -> Maybe BS.ByteString
decode :: ByteString -> Maybe ByteString
decode ByteString
bs
  | Bool
Arm.base64_arm_available = ByteString -> Maybe ByteString
Arm.decode ByteString
bs
  | Bool
otherwise                = ByteString -> Maybe ByteString
decode_scalar ByteString
bs
{-# INLINABLE decode #-}

encode_scalar :: BS.ByteString -> BS.ByteString
encode_scalar :: ByteString -> ByteString
encode_scalar (BI.PS ForeignPtr Word8
sfp Int
soff Int
l) =
  case ByteString
enc_tab of
    BI.PS ForeignPtr Word8
tfp Int
toff Int
_ ->
      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 ->
        ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
tfp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
tp0 -> 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
              !tp :: Ptr Word8
tp = Ptr Word8
tp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
toff :: Ptr Word8
              !nfull :: Int
nfull = Int
l Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
3
              !rmn :: Int
rmn   = Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
              loop :: Int -> IO ()
loop !Int
i
                | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nfull = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    let !ii :: Int
ii = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                        !oo :: Int
oo = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
                    b0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp Int
ii
                    b1 <- peekElemOff sp (ii + 1)
                    b2 <- peekElemOff sp (ii + 2)
                    c0 <- peekElemOff tp (fi (b0 `B.shiftR` 2))
                    c1 <- peekElemOff tp (fi
                            (((b0 .&. 0x03) `B.shiftL` 4)
                         .|.  (b1 `B.shiftR` 4)))
                    c2 <- peekElemOff tp (fi
                            (((b1 .&. 0x0F) `B.shiftL` 2)
                         .|.  (b2 `B.shiftR` 6)))
                    c3 <- peekElemOff tp (fi (b2 .&. 0x3F))
                    pokeElemOff dst  oo      (c0 :: Word8)
                    pokeElemOff dst (oo + 1) c1
                    pokeElemOff dst (oo + 2) c2
                    pokeElemOff dst (oo + 3) c3
                    loop (i + 1)
          Int -> IO ()
loop Int
0
          case Int
rmn of
            Int
0 -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Int
1 -> do
              let !ii :: Int
ii = Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                  !oo :: Int
oo = Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
              b0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp Int
ii
              c0 <- peekElemOff tp (fi (b0 `B.shiftR` 2))
              c1 <- peekElemOff tp (fi ((b0 .&. 0x03) `B.shiftL` 4))
              pokeElemOff dst  oo      (c0 :: Word8)
              pokeElemOff dst (oo + 1) c1
              pokeElemOff dst (oo + 2) 0x3D
              pokeElemOff dst (oo + 3) 0x3D
            Int
_ -> do
              let !ii :: Int
ii = Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                  !oo :: Int
oo = Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
              b0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp Int
ii
              b1 <- peekElemOff sp (ii + 1)
              c0 <- peekElemOff tp (fi (b0 `B.shiftR` 2))
              c1 <- peekElemOff tp (fi
                      (((b0 .&. 0x03) `B.shiftL` 4)
                   .|.  (b1 `B.shiftR` 4)))
              c2 <- peekElemOff tp (fi ((b1 .&. 0x0F) `B.shiftL` 2))
              pokeElemOff dst  oo      (c0 :: Word8)
              pokeElemOff dst (oo + 1) c1
              pokeElemOff dst (oo + 2) c2
              pokeElemOff dst (oo + 3) 0x3D

decode_scalar :: BS.ByteString -> Maybe BS.ByteString
decode_scalar :: ByteString -> Maybe ByteString
decode_scalar (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 = case ByteString
dec_tab of
      BI.PS ForeignPtr Word8
tfp Int
toff Int
_ -> 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 ->
        ForeignPtr Word8
-> (Ptr Word8 -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
tfp ((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
tp0 -> 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
              !tp :: Ptr Word8
tp = Ptr Word8
tp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
toff :: 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
                  !nbody  = if Int
pad Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Int
nfull Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 else Int
nfull
                  !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 -> do
                let body_loop :: Word8 -> Int -> IO Word8
body_loop !Word8
acc !Int
i
                      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nbody = Word8 -> IO Word8
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word8
acc
                      | Bool
otherwise = do
                          let !ii :: Int
ii = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2
                              !oo :: Int
oo = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                          c0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp  Int
ii
                          c1 <- peekElemOff sp (ii + 1)
                          c2 <- peekElemOff sp (ii + 2)
                          c3 <- peekElemOff sp (ii + 3)
                          v0 <- peekElemOff tp (fi c0)
                          v1 <- peekElemOff tp (fi c1)
                          v2 <- peekElemOff tp (fi c2)
                          v3 <- peekElemOff tp (fi c3)
                          let !b0 = (Word8
v0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
v1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
4) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03)
                              !b1 = ((Word8
v1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0F) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
4)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
v2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
2) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0F)
                              !b2 = ((Word8
v2 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
6)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. (Word8
v3 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x3F)
                          pokeElemOff dst  oo      b0
                          pokeElemOff dst (oo + 1) b1
                          pokeElemOff dst (oo + 2) b2
                          body_loop
                            (acc .|. v0 .|. v1 .|. v2 .|. v3) (i + 1)
                acc <- Word8 -> Int -> IO Word8
body_loop Word8
0 Int
0
                if acc .&. 0x80 /= 0
                  then pure False
                  else case pad of
                    Int
0 -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
                    Int
1 -> do
                      let !ii :: Int
ii = Int
nbody Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2
                          !oo :: Int
oo = Int
nbody Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                      c0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp  Int
ii
                      c1 <- peekElemOff sp (ii + 1)
                      c2 <- peekElemOff sp (ii + 2)
                      v0 <- peekElemOff tp (fi c0)
                      v1 <- peekElemOff tp (fi c1)
                      v2 <- peekElemOff tp (fi c2)
                      let !tail_acc = Word8
v0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
v1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
v2
                      if tail_acc .&. 0x80 /= 0 || v2 .&. 0x03 /= 0
                        then pure False
                        else do
                          let !b0 = (Word8
v0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
v1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
4) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03)
                              !b1 = ((Word8
v1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0F) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
4)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
v2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
2) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0F)
                          pokeElemOff dst  oo      b0
                          pokeElemOff dst (oo + 1) b1
                          pure True
                    Int
_ -> do
                      let !ii :: Int
ii = Int
nbody Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2
                          !oo :: Int
oo = Int
nbody Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
                      c0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp  Int
ii
                      c1 <- peekElemOff sp (ii + 1)
                      v0 <- peekElemOff tp (fi c0)
                      v1 <- peekElemOff tp (fi c1)
                      let !tail_acc = Word8
v0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
v1
                      if tail_acc .&. 0x80 /= 0 || v1 .&. 0x0F /= 0
                        then pure False
                        else do
                          let !b0 = (Word8
v0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
2)
                                Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. ((Word8
v1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
4) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x03)
                          pokeElemOff dst oo b0
                          pure True
              pure $! if ok then Just (BI.PS fp 0 outlen) else Nothing