content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Thrift.Protocol.JSON.Base64
    ( encodeBase64
    , encodeBase64Text
    , decodeBase64
    ) where
  
  import Data.Bits
  import Data.ByteString.Builder
  import Data.ByteString.Builder.Prim as Prim
  import Data.ByteString.Internal
  import Data.ByteString.Lazy (toStrict)
  import Data.Text (Text)
  import Data.Word
  import Data.Vector ((!))
  import qualified Data.ByteString as BS
  import qualified Data.Text as Text
  import qualified Data.Vector as Vector
  
  symbol :: Word8 -> Word8
  symbol w = BS.index
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    (fromIntegral w)
  
  encodeBase64 :: ByteString -> Builder
  encodeBase64 = encode . BS.unpack
    where
      -- case 1: len >= 3
      encode (b0 : b1 : b2 : bs) = primFixed
        (Prim.word8 >*< Prim.word8 >*< Prim.word8 >*< Prim.word8)
        (symbol $ (b0 `shiftR` 2) .&. 0x3f,
         (symbol $ ((b0 `shiftL` 4) .&. 0x30) .|. ((b1 `shiftR` 4) .&. 0x0f),
          (symbol $ ((b1 `shiftL` 2) .&. 0x3c) .|. ((b2 `shiftR` 6) .&. 0x03),
           symbol $ b2 .&. 0x3f))) <>
        encode bs
      -- case 2: len == 2
      encode [b0, b1] = primFixed
        (Prim.word8 >*< Prim.word8 >*< Prim.word8)
        (symbol $ (b0 `shiftR` 2) .&. 0x3f,
         (symbol $ ((b0 `shiftL` 4) .&. 0x30) .|. ((b1 `shiftR` 4) .&. 0x0f),
          symbol $ (b1 `shiftL` 2) .&. 0x3c))
      -- case 3: len == 1
      encode [b0] = primFixed
        (Prim.word8 >*< Prim.word8)
        (symbol $ (b0 `shiftR` 2) .&. 0x3f,
         symbol $ (b0 `shiftL` 4) .&. 0x30)
      -- case 4: len == 0 (impossible)
      encode [] = mempty
  
  encodeBase64Text :: ByteString -> Text
  encodeBase64Text =
    Text.pack . map w2c . BS.unpack . toStrict . toLazyByteString . encodeBase64
  
  decodeTable :: Word8 -> Word8
  decodeTable = (table!) . fromIntegral
    where
      table = Vector.fromList
        [ 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x3e, 0xff, 0xff, 0xff, 0x3f
        , 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06
        , 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12
        , 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x24
        , 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30
        , 0x31, 0x32, 0x33, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
        , 0xff, 0xff, 0xff, 0xff
        ]
  
  decodeBase64 :: ByteString -> ByteString
  decodeBase64 = BS.pack . decode . BS.unpack
    where
      decode (b0 : b1 : b2 : b3 : bs) =
        decode0 b0 b1 : decode1 b1 b2 : decode2 b2 b3 : decode bs
      decode [b0, b1, b2] = [decode0 b0 b1, decode1 b1 b2]
      decode [b0, b1] = [decode0 b0 b1]
      decode _ = []
  
      decode0 b0 b1 =
        (decodeTable b0 `shiftL` 2) .|.
        (decodeTable b1 `shiftR` 4)
      decode1 b1 b2 =
        ((decodeTable b1 `shiftL` 4) .&. 0xf0) .|.
        (decodeTable b2 `shiftR` 2)
      decode2 b2 b3 =
        ((decodeTable b2 `shiftL` 6) .&. 0xc0) .|.
        decodeTable b3
  