plum

module std

import std/Str

import std/Buffer

import std/Option

import std/Bool

import std/Number

# The Base64 package contains support for doing Base64 binary-to-text encodings.

#

# Implemented on top of `[]Byte`/`Buffer` (`libs/std/bytes.plum`/`buffer.plum`)

# — the original note here said a real bit-twiddling implementation couldn't

# be expressed without indexed byte-buffer mutation, which didn't exist in

# this language yet at the time it was written; it does now.

#

# `data`/the decoded payload are represented as `Str` throughout (matching

# the original public shape) rather than `[]Byte` — `Str` is itself just a

# byte array here (see `libs/std/str.plum`), so this is a plain byte string,

# not a claim that the content is valid UTF-8 text.

# Encode for PEM (RFC 1421): 64-character lines.

fun encodePEM(data: Str) -> Str =

  return encode(data, "+", "/", "=", 64, "\r\n")

# Encode for MIME (RFC 2045): 76-character lines.

fun encodeMIME(data: Str) -> Str =

  return encode(data, "+", "/", "=", 76, "\r\n")

# Encode for URLs (RFC 4648). Padding characters are stripped by default.

fun encodeURL(data: Str, pad: Bool = False) -> Str =

  return encode(data, "-", "_", pad ? "=" : "", 0, "")

# Configurable encoding. The defaults are for RFC 4648. `pad` is the padding

# character, or `""` to omit padding entirely; `linelen` is the output line

# length to wrap at (`0` disables wrapping), separated by `linesep`.

fun encode(data: Str, at62: Str = "+", at63: Str = "/", pad: Str = "=", linelen: Int = 0, linesep: Str = "\r\n") -> Str =

  at62b := at62.byteAt(0)

  at63b := at63.byteAt(0)

  has_pad := pad.length() > 0

  padb := has_pad ? pad.byteAt(0) : 0

  res := Buffer()

  n := data.length()

  i := 0

  line_chars := 0

  while i < n

    b0 := data.byteAt(i)

    has1 := i + 1 < n

    has2 := i + 2 < n

    b1 := has1 ? data.byteAt(i + 1) : 0

    b2 := has2 ? data.byteAt(i + 2) : 0

    res.writeByte(Byte(encodeByte(b0 >> 2, at62b, at63b)))

    res.writeByte(Byte(encodeByte({b0 & 0x03} << 4 | b1 >> 4, at62b, at63b)))

    line_chars = line_chars + 2

    if has1

      res.writeByte(Byte(encodeByte({b1 & 0x0F} << 2 | b2 >> 6, at62b, at63b)))

      line_chars = line_chars + 1

    else if has_pad

      res.writeByte(Byte(padb))

      line_chars = line_chars + 1

    if has2

      res.writeByte(Byte(encodeByte(b2 & 0x3F, at62b, at63b)))

      line_chars = line_chars + 1

    else if has_pad

      res.writeByte(Byte(padb))

      line_chars = line_chars + 1

    i = i + 3

    if linelen > 0 && line_chars >= linelen && i < n

      res.write(linesep)

      line_chars = 0

  return res.toStr()

# Encodes a single 6-bit value (0-63) as its base64 alphabet character.

fun encodeByte(i: Int, at62: Int, at63: Int) -> Int =

  if i < 26

    return 65 + i # 'A'..'Z'

  if i < 52

    return 97 + {i - 26} # 'a'..'z'

  if i < 62

    return 48 + {i - 52} # '0'..'9'

  if i == 62

    return at62

  return at63

# The inverse of `encodeByte`: a base64 alphabet character's 6-bit value, or

# -1 if `b` isn't one of the 64 alphabet characters (whitespace/padding are

# handled by the caller, not here).

fun decodeByte(b: Int, at62: Int, at63: Int) -> Int =

  if b >= 65 && b <= 90

    return b - 65

  if b >= 97 && b <= 122

    return b - 97 + 26

  if b >= 48 && b <= 57

    return b - 48 + 52

  if b == at62

    return 62

  if b == at63

    return 63

  return -1

fun isBase64Whitespace(b: Int) -> Bool =

  b == 32 || b == 9 || b == 10 || b == 13

# Decode for URLs (RFC 4648).

fun decodeUrl(data: Str) -> Option[Str] =

  return decode(data, "-", "_", "=")

# Configurable decoding. The defaults are for RFC 4648. Missing padding is

# not an error. Non-base64 data, other than whitespace (which can appear at

# any point), is an error (`None`).

fun decode(data: Str, at62: Str = "+", at63: Str = "/", pad: Str = "=") -> Option[Str] =

  at62b := at62.byteAt(0)

  at63b := at63.byteAt(0)

  has_pad := pad.length() > 0

  padb := has_pad ? pad.byteAt(0) : -1

  res := Buffer()

  n := data.length()

  i := 0

  bits := 0

  nbits := 0

  while i < n

    b := data.byteAt(i)

    i = i + 1

    if isBase64Whitespace(b) || b == padb

      continue

    v := decodeByte(b, at62b, at63b)

    if v < 0

      return None

    bits = {bits << 6} | v

    nbits = nbits + 6

    if nbits >= 8

      nbits = nbits - 8

      res.writeByte(Byte({bits >> nbits} & 0xFF))

  return Some(res.toStr())