base64faster.nim

import strutils


proc initDecodeTable*(): array[256, char] =
  # computes a decode table at compile time
  for i in 0 ..< 256:
    let ch = char(i)
    var code = 255
    if ch >= 'A' and ch <= 'Z': code = i - 0x00000041
    if ch >= 'a' and ch <= 'z': code = i - 0x00000047
    if ch >= '0' and ch <= '9': code = i + 0x00000004
    if ch == '+' or ch == '-': code = 0x0000003E
    if ch == '/' or ch == '_': code = 0x0000003F
    result[i] = char(code)


const
  chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
  decodeTable = initDecodeTable()


proc decode*(str: string): string =
  ## Decodes a string in base64 representation back into its original form.

  proc decodeSize(size: int): int =
    return (size * 3 div 4) + 6

  template inputChar(x: untyped) =
    let x = int decode_table[ord(str[inputIndex])]
    inc inputIndex
    if x == 255:
      echo repr(result)
      raise newException(ValueError, "Invalid base64 format")

  template outputChar(x: untyped) =
    result[outputIndex] = char(x and 255)
    inc outputIndex

  # pre allocate output string once
  result.setLen(decodeSize(str.len))
  var
    inputIndex = 0
    outputIndex = 0
    inputLen = str.len
    inputEnds = 0
  # strip trailing characters
  while str[inputLen - 1] in  {'\n', '\r', '='}:
    dec inputLen
  # read 4 characters at at time
  inputEnds = inputLen - 4
  while inputIndex <= inputEnds:
    while str[inputIndex] in {'\n', '\r'}:
      inc inputIndex
    inputChar(a)
    inputChar(b)
    inputChar(c)
    inputChar(d)
    outputChar(a shl 2 or b shr 4)
    outputChar(b shl 4 or c shr 2)
    outputChar(c shl 6 or d shr 0)
  # do the last 2 or 3 characters
  var leftLen = abs((inputIndex - inputLen) mod 4)
  if leftLen == 2:
    inputChar(a)
    inputChar(b)
    outputChar(a shl 2 or b shr 4)
  elif leftLen == 3:
    inputChar(a)
    inputChar(b)
    inputChar(c)
    outputChar(a shl 2 or b shr 4)
    outputChar(b shl 4 or c shr 2)
  result.setLen(outputIndex)


proc encode*(str: string): string =
  ## Encodes `s` into base64 representation.

  proc encodeSize(size: int): int =
    return (size * 4 div 3) + 6

  template inputByteFirst(shift: int) =
    n = uint32(str[inputIndex]) shl shift
    inc inputIndex

  template inputByteNext(shift: int) =
    n = n or uint32(str[inputIndex]) shl shift
    inc inputIndex

  template outputChar(x: untyped) =
    result[outputIndex] = chars[x and 63]
    inc outputIndex

  template outputChar(c: char) =
    result[outputIndex] = c
    inc outputIndex

  result.setLen(encodeSize(str.len))
  var
    inputIndex = 0
    outputIndex = 0
    inputEnds = str.len - str.len mod 3
    n: uint32

  while inputIndex != inputEnds:
    inputByteFirst(16)
    inputByteNext(8)
    inputByteNext(0)
    outputChar(n shr 18)
    outputChar(n shr 12)
    outputChar(n shr 6)
    outputChar(n shr 0)

  var padding = str.len mod 3
  if padding == 1:
    inputByteFirst(16)
    outputChar(n shr 18)
    outputChar(n shr 12)
    outputChar('=')
    outputChar('=')

  elif padding == 2:
    inputByteFirst(16)
    inputByteNext(8)
    outputChar(n shr 18)
    outputChar(n shr 12)
    outputChar(n shr 6)
    outputChar('=')

  result.setLen(outputIndex)


when isMainModule:
  import times

  let STR_SIZE = 10_000_000
  let TRIES = 100
  let str = strutils.repeat('a', STR_SIZE)
  var str2 = ""

  block:
    var t = times.epochTime()
    var i = 0
    var s: int64 = 0
    while i < TRIES:
      # set large line length to avoid adding line breaks,
      # outputing identical encoded strings as in other languages
      str2 = encode(str)
      s += len(str2)
      i += 1
    echo("encode: ", s, ", ", formatFloat(times.epochTime() - t, ffDefault, 6))

  block:
    var t = times.epochTime()
    var i = 0
    var s: int64 = 0
    while i < TRIES:
      s += len(decode(str2))
      i += 1
    echo("decode: ", s, ", ", formatFloat(times.epochTime() - t, ffDefault, 6))