Bech32.scala

package scoin

import scodec.bits.ByteVector

/** Bech32 and Bech32m address formats. See
  * https://github.com/bitcoin/bips/blob/master/bip-0173.mediawiki and
  * https://github.com/bitcoin/bips/blob/master/bip-0350.mediawiki.
  */
object Bech32 {
  val alphabet = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

  // @formatter:off
  sealed trait Encoding
  case object Bech32Encoding extends Encoding
  case object Bech32mEncoding extends Encoding
  // @formatter:on

  // 5 bits integer
  // Bech32 works with 5 bits values, we use this type to make it explicit: whenever you see Int5 it means 5 bits values,
  // and whenever you see Byte it means 8 bits values
  type Int5 = Byte

  // char -> 5 bits value
  private val InvalidChar = 255.toByte
  val map = {
    val result = new Array[Int5](255)
    for (i <- result.indices) result(i) = InvalidChar
    alphabet.zipWithIndex.foreach { case (c, i) => result(c) = i.toByte }
    result
  }

  private def expand(hrp: String): Array[Int5] = {
    val result = new Array[Int5](2 * hrp.length + 1)
    var i = 0
    while (i < hrp.length) {
      result(i) = (hrp(i).toInt >>> 5).toByte
      result(hrp.length() + 1 + i) = (hrp(i).toInt & 31).toByte
      i = i + 1
    }
    result(hrp.length()) = 0.toByte
    result
  }

  private def polymod(values: Array[Int5], values1: Array[Int5]): Int = {
    val GEN = Array(0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3)
    var chk = 1
    values.foreach(v => {
      val b = chk >>> 25
      chk = ((chk & 0x1ffffff) << 5) ^ v
      for (i <- 0 until 5) {
        if (((b >>> i) & 1) != 0) chk = chk ^ GEN(i)
      }
    })
    values1.foreach(v => {
      val b = chk >>> 25
      chk = ((chk & 0x1ffffff) << 5) ^ v
      for (i <- 0 until 5) {
        if (((b >>> i) & 1) != 0) chk = chk ^ GEN(i)
      }
    })
    chk
  }

  /** @param hrp
    *   human readable prefix
    * @param int5s
    *   5-bit data
    * @return
    *   hrp + data encoded as a Bech32 string
    */
  def encode(hrp: String, int5s: Array[Int5], encoding: Encoding): String = {
    require(
      hrp.toLowerCase == hrp || hrp.toUpperCase == hrp,
      "mixed case strings are not valid bech32 prefixes"
    )
    val checksum = Bech32.checksum(hrp, int5s, encoding)
    hrp + "1" + new String((int5s ++ checksum).map(i => alphabet(i)))
  }

  /** decodes a bech32 or bech32m string
    *
    * @param bech32
    *   bech32 or bech32m string
    * @return
    *   a (encoding, hrp, data) tuple
    */
  def decode(bech32: String): (String, Array[Int5], Encoding) = {
    require(
      bech32.toLowerCase == bech32 || bech32.toUpperCase == bech32,
      "mixed case strings are not valid bech32"
    )
    bech32.foreach(c => require(c >= 33 && c <= 126, "invalid character"))
    val input = bech32.toLowerCase()
    val pos = input.lastIndexOf('1')
    val hrp = input.take(pos)
    require(
      hrp.nonEmpty && hrp.length <= 83,
      "hrp must contain 1 to 83 characters"
    )
    val data = new Array[Int5](input.length - pos - 1)
    for (i <- data.indices) {
      val elt = map(input(pos + 1 + i))
      require(
        elt != InvalidChar,
        s"invalid bech32 character ${input(pos + 1 + i)}"
      )
      data(i) = elt
    }
    val checksum = polymod(expand(hrp), data)
    val encoding = checksum match {
      case 1          => Bech32Encoding
      case 0x2bc830a3 => Bech32mEncoding
      case _ =>
        throw new IllegalArgumentException(s"invalid checksum for $bech32")
    }
    (hrp, data.dropRight(6), encoding)
  }

  /** @param hrp
    *   Human Readable Part
    * @param data
    *   data (a sequence of 5 bits integers)
    * @param encoding
    *   encoding to use (bech32 or bech32m)
    * @return
    *   a checksum computed over hrp and data
    */
  private def checksum(
      hrp: String,
      data: Array[Int5],
      encoding: Encoding
  ): Array[Int5] = {
    val constant = encoding match {
      case Bech32Encoding  => 1
      case Bech32mEncoding => 0x2bc830a3
    }
    val values = expand(hrp) ++ data
    val poly = polymod(
      values,
      Array(0.toByte, 0.toByte, 0.toByte, 0.toByte, 0.toByte, 0.toByte)
    ) ^ constant
    val result = new Array[Int5](6)
    for (i <- 0 to 5) result(i) = ((poly >>> 5 * (5 - i)) & 31).toByte
    result
  }

  /** @param input
    *   a sequence of 8 bits integers
    * @return
    *   a sequence of 5 bits integers
    */
  def eight2five(input: Array[Byte]): Array[Int5] = {
    var buffer = 0L
    val output = collection.mutable.ArrayBuffer.empty[Byte]
    var count = 0
    input.foreach(b => {
      buffer = (buffer << 8) | (b & 0xff)
      count = count + 8
      while (count >= 5) {
        output.append(((buffer >> (count - 5)) & 31).toByte)
        count = count - 5
      }
    })
    if (count > 0) output.append(((buffer << (5 - count)) & 31).toByte)
    output.toArray
  }

  /** @param input
    *   a sequence of 5 bits integers
    * @return
    *   a sequence of 8 bits integers
    */
  def five2eight(input: Array[Int5]): Array[Byte] = {
    var buffer = 0L
    val output = collection.mutable.ArrayBuffer.empty[Byte]
    var count = 0
    input.foreach(b => {
      buffer = (buffer << 5) | (b & 31)
      count = count + 5
      while (count >= 8) {
        output.append(((buffer >> (count - 8)) & 0xff).toByte)
        count = count - 8
      }
    })
    require(count <= 4, "Zero-padding of more than 4 bits")
    require(
      (buffer & ((1 << count) - 1)) == 0,
      "Non-zero padding in 8-to-5 conversion"
    )
    output.toArray
  }

  /** encode a bitcoin witness address
    *
    * @param hrp
    *   should be "bc" or "tb"
    * @param witnessVersion
    *   witness version (0 to 16)
    * @param data
    *   witness program: if version is 0, either 20 bytes (P2WPKH) or 32 bytes
    *   (P2WSH)
    * @return
    *   a bech32 encoded witness address
    */
  def encodeWitnessAddress(
      hrp: String,
      witnessVersion: Byte,
      data: ByteVector
  ): String = {
    require(
      0 <= witnessVersion && witnessVersion <= 16,
      "invalid segwit version"
    )
    val encoding = witnessVersion match {
      case 0 => Bech32Encoding
      case _ => Bech32mEncoding
    }
    val data1 = witnessVersion +: Bech32.eight2five(data.toArray)
    val checksum = Bech32.checksum(hrp, data1, encoding)
    hrp + "1" + new String((data1 ++ checksum).map(i => alphabet(i)))
  }

  /** decode a bitcoin witness address
    *
    * @param address
    *   witness address
    * @return
    *   a (prefix, version, program) tuple where prefix is the human-readable
    *   prefix, version is the witness version and program the decoded witness
    *   program. If version is 0, it will be either 20 bytes (P2WPKH) or 32
    *   bytes (P2WSH).
    */
  def decodeWitnessAddress(address: String): (String, Byte, ByteVector) = {
    if (
      address.indexWhere(_.isLower) != -1 && address.indexWhere(_.isUpper) != -1
    )
      throw new IllegalArgumentException(
        "input mixes lowercase and uppercase characters"
      )
    val (hrp, data, encoding) = decode(address)
    require(hrp == "bc" || hrp == "tb" || hrp == "bcrt", s"invalid HRP $hrp")
    val version = data(0)
    require(version >= 0 && version <= 16, "invalid segwit version")
    val bin = five2eight(data.drop(1))
    require(
      bin.length >= 2 && bin.length <= 40,
      s"invalid witness program length ${bin.length}"
    )
    if (version == 0)
      require(
        encoding == Bech32Encoding,
        "version 0 must be encoded with Bech32"
      )
    if (version == 0)
      require(
        bin.length == 20 || bin.length == 32,
        s"invalid witness program length ${bin.length}"
      )
    if (version != 0)
      require(
        encoding == Bech32mEncoding,
        "version 1 to 16 must be encoded with Bech32m"
      )
    (hrp, version, ByteVector.view(bin))
  }
}