{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE NoStarIsType #-}

module Data.Bytes.Base64
  ( encode
  , builder
  , recodeBoundedBuilder
  , decode64
  ) where

import Control.Monad.ST.Run (runByteArrayST)
import Data.Bits (unsafeShiftL, unsafeShiftR, (.&.), (.|.))
import Data.Bytes.Types (Bytes (Bytes))
import Data.Char (ord)
import Data.Primitive (ByteArray (..), MutableByteArray (..), newByteArray, readByteArray, unsafeFreezeByteArray)
import Data.Primitive.Ptr (indexOffPtr)
import Data.Word (Word64, Word8)
import GHC.Exts (Int (I#), Ptr (Ptr), State#, (+#), (-#))
import GHC.ST (ST (ST))
import GHC.TypeNats (Div, type (*), type (+))
import GHC.Word (Word (W#), Word32)
import GHC.Word.Compat (pattern W32#)

import qualified Arithmetic.Nat as Nat
import qualified Arithmetic.Types as Arithmetic
import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Builder.Bounded.Unsafe as BBU
import qualified Data.Bytes.Builder.Unsafe as BU
import qualified Data.Primitive.ByteArray.BigEndian as BE
import qualified Data.Primitive.ByteArray.LittleEndian as LE
import qualified Data.Primitive.Ptr as PM
import qualified GHC.Exts as Exts

-- | Encode a byte sequence with base64.
encode :: Bytes -> ByteArray
encode :: Bytes -> ByteArray
encode (Bytes ByteArray
src Int
soff Int
slen) = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST do
  let dlen :: Int
dlen = Int -> Int
calculatePaddedLength Int
slen
  MutableByteArray s
dst <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: Type -> Type).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
dlen
  MutableByteArray s -> Int -> ByteArray -> Int -> Int -> ST s ()
forall s.
MutableByteArray s -> Int -> ByteArray -> Int -> Int -> ST s ()
performEncodeImmutable MutableByteArray s
dst Int
0 ByteArray
src Int
soff Int
slen
  MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: Type -> Type).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst

-- | Encode a byte sequence with base64 as a builder.
builder :: Bytes -> BU.Builder
builder :: Bytes -> Builder
builder (Bytes ByteArray
src Int
soff Int
slen) = Int -> (forall s. MutableByteArray s -> Int -> ST s Int) -> Builder
BU.fromEffect Int
dlen \MutableByteArray s
dst Int
doff -> do
  MutableByteArray s -> Int -> ByteArray -> Int -> Int -> ST s ()
forall s.
MutableByteArray s -> Int -> ByteArray -> Int -> Int -> ST s ()
performEncodeImmutable MutableByteArray s
dst Int
doff ByteArray
src Int
soff Int
slen
  Int -> ST s Int
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Int
doff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dlen)
 where
  dlen :: Int
dlen = Int -> Int
calculatePaddedLength Int
slen

{- | Encode a byte sequence with base64. This augments a builder
by first playing the original builder to paste the byte sequence
and then performing the base64 encoding in-place on the buffer
that had been pasted into.
-}
recodeBoundedBuilder ::
  Arithmetic.Nat n ->
  BBU.Builder n ->
  BBU.Builder (4 * (Div (n + 2) 3))
recodeBoundedBuilder :: forall (n :: Nat).
Nat n -> Builder n -> Builder (4 * Div (n + 2) 3)
recodeBoundedBuilder !Nat n
n (BBU.Builder forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, Int# #)
f) =
  (forall s.
 MutableByteArray# s -> Int# -> State# s -> (# State# s, Int# #))
-> Builder (4 * Div (n + 2) 3)
forall (a :: Nat).
(forall s.
 MutableByteArray# s -> Int# -> State# s -> (# State# s, Int# #))
-> Builder a
BBU.Builder
    ( \MutableByteArray# s
arr Int#
off0 State# s
s0 ->
        let !off1 :: Int#
off1 = (Int#
off0 Int# -> Int# -> Int#
+# Int#
maxEncLen) Int# -> Int# -> Int#
-# Int#
maxRawLen
         in case MutableByteArray# s -> Int# -> State# s -> (# State# s, Int# #)
forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, Int# #)
f MutableByteArray# s
arr Int#
off1 State# s
s0 of
              (# State# s
s1, Int#
off2 #) ->
                let !actualLen :: Int#
actualLen = Int#
off2 Int# -> Int# -> Int#
-# Int#
off1
                 in case ST s () -> State# s -> (# State# s, () #)
forall s a. ST s a -> State# s -> (# State# s, a #)
unST (MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
forall s.
MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
performEncode (MutableByteArray# s -> MutableByteArray s
forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
arr) (Int# -> Int
I# Int#
off0) (MutableByteArray# s -> MutableByteArray s
forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
arr) (Int# -> Int
I# Int#
off1) (Word# -> Word
W# (Int# -> Word#
Exts.int2Word# Int#
actualLen))) State# s
s1 of
                      (# State# s
s2, (()
_ :: ()) #) ->
                        let !(I# Int#
actualEncLen) = Int -> Int
calculatePaddedLength (Int# -> Int
I# Int#
actualLen)
                         in (# State# s
s2, Int#
actualEncLen #)
    )
 where
  !(I# Int#
maxRawLen) = Nat n -> Int
forall (n :: Nat). Nat n -> Int
Nat.demote Nat n
n
  !(I# Int#
maxEncLen) = Int -> Int
calculatePaddedLength (Int# -> Int
I# Int#
maxRawLen)

performEncodeImmutable ::
  MutableByteArray s -> -- dest
  Int -> -- dest off
  ByteArray ->
  Int -> -- src off
  Int -> -- source length
  ST s ()
performEncodeImmutable :: forall s.
MutableByteArray s -> Int -> ByteArray -> Int -> Int -> ST s ()
performEncodeImmutable MutableByteArray s
dst Int
doff (ByteArray ByteArray#
src) Int
soff Int
slen =
  MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
forall s.
MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
performEncode MutableByteArray s
dst Int
doff (MutableByteArray# s -> MutableByteArray s
forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray (ByteArray# -> MutableByteArray# s
forall a b. a -> b
Exts.unsafeCoerce# ByteArray#
src)) Int
soff (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word Int
slen)

-- The function is the meat of this module. Implementation notes:
--
-- We use big-endian and little-endian unaligned loads and stores
-- from the byte-order library. This means we can cut down loads
-- and stores by a factor of 4.
--
-- Once we get down to less than 4 bytes, we have to do some grunt
-- work to finish off the encoding. This happens for two reasons.
-- The first is that base64 requires trailing equals signs to pad
-- encoded byte sequences whose length was not a multiple of three
-- The second is that our 32-bit load is unsafe once we are at the
-- end since its possible (although really unlikely) that the byte
-- sequences is right up against the end of the address space that
-- is available to GHC. Segfaults happen if we wander outside of
-- this pasture.
--
-- Why is the source a mutable byte array? We actually need this
-- to accept both immutable and mutable byte arrays as the source.
-- To avoid code duplication, we use unsafeCoerce# in performEncodeImmutable.
-- Using the mutable variant here actually gives us slightly better
-- guarantees from the compiler since read (unlike index) is sequenced.
-- These guarantees are important in recodeBoundedBuilder, where the
-- encoding is performed in-place.
--
-- Also, what's the deal with the source length being a Word instead
-- of an Int. GHC can actually generate code when we do this.
-- In the cmm stage of compilation, case-on-number constructs
-- are compiled to lower-level constructs. They become either jump
-- table or a series of conditionals statements. In our case,
-- an unsigned number helps GHC realize that it does not need
-- to test for n<0, although it still must test for n>3.
performEncode ::
  MutableByteArray s -> -- dest
  Int -> -- dest off
  MutableByteArray s -> -- src
  Int -> -- src off
  Word -> -- source length
  ST s ()
performEncode :: forall s.
MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
performEncode !MutableByteArray s
dst !Int
doff !MutableByteArray s
src !Int
soff !Word
slen = case Word
slen of
  Word
3 -> do
    Word8
x1 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src Int
soff
    Word8
x2 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src (Int
soff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Word8
x3 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src (Int
soff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
    let (Word
w1, Word
w2, Word
w3, Word
w4) = Word32 -> (Word, Word, Word, Word)
disassembleBE (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleBE Word8
x1 Word8
x2 Word8
x3 Word8
0)
        c1 :: Word8
c1 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w1)
        c2 :: Word8
c2 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w2)
        c3 :: Word8
c3 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w3)
        c4 :: Word8
c4 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w4)
    MutableByteArray (PrimState (ST s)) -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, PrimUnaligned a, Bytes a) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
LE.writeUnalignedByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
doff (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE Word8
c1 Word8
c2 Word8
c3 Word8
c4)
  Word
2 -> do
    Word8
x1 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src Int
soff
    Word8
x2 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src (Int
soff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    let (Word
w1, Word
w2, Word
w3, Word
_) = Word32 -> (Word, Word, Word, Word)
disassembleBE (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleBE Word8
x1 Word8
x2 Word8
0 Word8
0)
        c1 :: Word8
c1 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w1)
        c2 :: Word8
c2 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w2)
        c3 :: Word8
c3 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w3)
        c4 :: Word8
c4 = Char -> Word8
c2w Char
'='
    MutableByteArray (PrimState (ST s)) -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, PrimUnaligned a, Bytes a) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
LE.writeUnalignedByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
doff (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE Word8
c1 Word8
c2 Word8
c3 Word8
c4)
  Word
1 -> do
    Word8
x1 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src Int
soff
    let (Word
w1, Word
w2, Word
_, Word
_) = Word32 -> (Word, Word, Word, Word)
disassembleBE (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleBE Word8
x1 Word8
0 Word8
0 Word8
0)
        c1 :: Word8
c1 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w1)
        c2 :: Word8
c2 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w2)
        c3 :: Word8
c3 = Char -> Word8
c2w Char
'='
        c4 :: Word8
c4 = Char -> Word8
c2w Char
'='
    MutableByteArray (PrimState (ST s)) -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, PrimUnaligned a, Bytes a) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
LE.writeUnalignedByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
doff (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE Word8
c1 Word8
c2 Word8
c3 Word8
c4)
  Word
0 -> () -> ST s ()
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()
  Word
_ -> do
    -- This last case is always slen > 3
    Word32
w :: Word32 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word32
forall (m :: Type -> Type) a.
(PrimMonad m, PrimUnaligned a, Bytes a) =>
MutableByteArray (PrimState m) -> Int -> m a
BE.readUnalignedByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
src Int
soff
    let (Word
w1, Word
w2, Word
w3, Word
w4) = Word32 -> (Word, Word, Word, Word)
disassembleBE Word32
w
        c1 :: Word8
c1 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w1)
        c2 :: Word8
c2 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w2)
        c3 :: Word8
c3 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w3)
        c4 :: Word8
c4 = Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
indexOffPtr Ptr Word8
table (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int Word
w4)
    MutableByteArray (PrimState (ST s)) -> Int -> Word32 -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, PrimUnaligned a, Bytes a) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
LE.writeUnalignedByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
doff (Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE Word8
c1 Word8
c2 Word8
c3 Word8
c4)
    MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
forall s.
MutableByteArray s
-> Int -> MutableByteArray s -> Int -> Word -> ST s ()
performEncode MutableByteArray s
dst (Int
doff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) MutableByteArray s
src (Int
soff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) (Word
slen Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
3)

-- Argument bytes are hi to lo. The first argument becomes
-- the least significant component of the result.
assembleLE :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleLE Word8
a Word8
b Word8
c Word8
d =
  Word -> Word32
unsafeW32
    ( Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
d) Int
24
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
c) Int
16
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
b) Int
8
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
a)
    )

-- Argument bytes are hi to lo. The first argument becomes
-- the most significant component of the result.
assembleBE :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleBE :: Word8 -> Word8 -> Word8 -> Word8 -> Word32
assembleBE Word8
a Word8
b Word8
c Word8
d =
  Word -> Word32
unsafeW32
    ( Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
a) Int
24
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
b) Int
16
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftL (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
c) Int
8
        Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word Word8
d)
    )

-- Create a 32-bit word from a machine word that we know
-- is not too large.
unsafeW32 :: Word -> Word32
unsafeW32 :: Word -> Word32
unsafeW32 (W# Word#
w) = Word# -> Word32
W32# Word#
w

-- We only care about the upper 24 bits of the argument.
-- This gets broken into four 6-bit words.
disassembleBE :: Word32 -> (Word, Word, Word, Word)
disassembleBE :: Word32 -> (Word, Word, Word, Word)
disassembleBE !Word32
w =
  ( Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word Word32
w) Int
26
  , Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word Word32
w) Int
20 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
0b00111111
  , Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word Word32
w) Int
14 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
0b00111111
  , Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word Word32
w) Int
8 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
0b00111111
  )

table :: Ptr Word8
table :: Ptr Word8
table = Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"#

calculatePaddedLength :: Int -> Int
calculatePaddedLength :: Int -> Int
calculatePaddedLength Int
n = Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int -> Int -> Int
divRoundUp Int
n Int
3)

divRoundUp :: Int -> Int -> Int
divRoundUp :: Int -> Int -> Int
divRoundUp Int
x Int
y = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
y

c2w :: Char -> Word8
c2w :: Char -> Word8
c2w = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Char -> Int) -> Char -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord

unST :: ST s a -> State# s -> (# State# s, a #)
unST :: forall s a. ST s a -> State# s -> (# State# s, a #)
unST (ST STRep s a
f) State# s
s = STRep s a
f State# s
s

-- Decode a base64-url-encoded 64-bit word. Rejects encoded numbers
-- greater than or equal to @2^64@. This maps the rightmost byte to
-- the 6 least significant bits of the word.
decode64 :: Bytes -> Maybe Word64
decode64 :: Bytes -> Maybe Word64
decode64 Bytes
bs
  | Bytes -> Int
Bytes.length Bytes
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10 = Maybe Word64
forall a. Maybe a
Nothing
  | Bool
otherwise =
      (Word64 -> Word8 -> Maybe Word64)
-> Word64 -> Bytes -> Maybe Word64
forall (m :: Type -> Type) a.
Monad m =>
(a -> Word8 -> m a) -> a -> Bytes -> m a
Bytes.foldlM
        ( \ !(Word64
acc :: Word64) Word8
b -> case Ptr Word8 -> Int -> Word8
forall a. Prim a => Ptr a -> Int -> a
PM.indexOffPtr Ptr Word8
decodeTable (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Int Word8
b) of
            Word8
0xFF -> Maybe Word64
forall a. Maybe a
Nothing
            Word8
w -> Word64 -> Maybe Word64
forall a. a -> Maybe a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Word64 -> Maybe Word64) -> Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$! (Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftL Word64
acc Int
6 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Word64 Word8
w)
        )
        Word64
0
        Bytes
bs

decodeTable :: Ptr Word8
decodeTable :: Ptr Word8
decodeTable =
  Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr
    Addr#
"\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\x3e\xff\xff\xff\x3f\
    \\x34\x35\x36\x37\x38\x39\x3a\x3b\x3c\x3d\xff\xff\xff\x63\xff\xff\
    \\xff\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\
    \\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\xff\xff\xff\xff\xff\
    \\xff\x1a\x1b\x1c\x1d\x1e\x1f\x20\x21\x22\x23\x24\x25\x26\x27\x28\
    \\x29\x2a\x2b\x2c\x2d\x2e\x2f\x30\x31\x32\x33\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
    \\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"#