Base32 / Base58 / Base62 / Base85 Encoder — Go source

// Package baseencoder is the Go twin of CosmoDev's src/lib/base-encoder.ts
// (dual source: the web lib is TypeScript, the CLI lib is Go — kept in
// lock-step). Pure + deterministic, never panics. The table-driven tests in
// base-encoder_test.go share vectors with src/lib/base-encoder.test.ts so the
// two implementations are held to the same contract.
//
// All four schemes operate on the UTF-8 bytes of the input text, mirroring the
// TS lib exactly: Base32 (RFC 4648, padded), Base58 (Bitcoin, leading-zero
// preserving), Base62 (big.Int base-conversion, no leading-zero special-casing),
// and Base85 (Ascii85, 'z' shorthand for full zero groups).
package baseencoder

import (
	"errors"
	"math/big"
	"strings"
)

// Scheme selects a byte-array base encoding. It mirrors the TS `Scheme` union
// ('base32' | 'base58' | 'base62' | 'base85').
type Scheme string

const (
	Base32 Scheme = "base32"
	Base58 Scheme = "base58"
	Base62 Scheme = "base62"
	Base85 Scheme = "base85"
)

// errInvalid is returned when an encoded string contains a character outside
// the scheme's alphabet or is otherwise malformed. It mirrors the TS libs'
// `null` return from the internal decode helpers.
var errInvalid = errors.New("baseencoder: invalid or malformed input")

const (
	b32Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
	b58Alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
	b62Alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
)

// outLen32 maps the size of a final (partial) 5-byte chunk to the number of
// data characters it emits before '=' padding, per RFC 4648. Index = byte
// count (1..4). Matches `outLen = [0, 2, 4, 5, 7][chunk.length]` in the TS.
var outLen32 = [5]int{0, 2, 4, 5, 7}

// ---------------------------------------------------------------------------
// Base32 — RFC 4648 alphabet, padded to a multiple of 8 chars with '='.
// ---------------------------------------------------------------------------

func encode32(data []byte) string {
	var sb strings.Builder
	for i := 0; i < len(data); i += 5 {
		end := i + 5
		if end > len(data) {
			end = len(data)
		}
		chunk := data[i:end]
		var b [5]int
		for j := 0; j < len(chunk); j++ {
			b[j] = int(chunk[j])
		}
		// Pack 5 bytes (40 bits) into 8 base32 digits (5 bits each, big-endian).
		digits := [8]int{
			(b[0] >> 3) & 0x1f,
			((b[0] << 2) | (b[1] >> 6)) & 0x1f,
			(b[1] >> 1) & 0x1f,
			((b[1] << 4) | (b[2] >> 4)) & 0x1f,
			((b[2] << 1) | (b[3] >> 7)) & 0x1f,
			(b[3] >> 2) & 0x1f,
			((b[3] << 3) | (b[4] >> 5)) & 0x1f,
			b[4] & 0x1f,
		}
		var group [8]byte
		for k, d := range digits {
			group[k] = b32Alphabet[d]
		}
		if len(chunk) == 5 {
			sb.Write(group[:])
		} else {
			n := outLen32[len(chunk)]
			sb.Write(group[:n])
			for p := 0; p < 8-n; p++ {
				sb.WriteByte('=')
			}
		}
	}
	return sb.String()
}

func decode32(s string) ([]byte, error) {
	var out []byte
	buffer := 0
	bits := 0
	for _, c := range s {
		if c == '=' {
			break // padding marks the end
		}
		idx := strings.IndexRune(b32Alphabet, c)
		if idx == -1 {
			return nil, errInvalid
		}
		buffer = (buffer << 5) | idx
		bits += 5
		if bits >= 8 {
			bits -= 8
			out = append(out, byte((buffer>>bits)&0xff))
			buffer &= (1 << bits) - 1 // keep only the leftover bits
		}
	}
	return out, nil
}

// ---------------------------------------------------------------------------
// Base58 — Bitcoin alphabet. Leading 0x00 bytes → leading '1' (count preserved).
// ---------------------------------------------------------------------------

func encode58(data []byte) string {
	// Count leading zero bytes — each maps to a leading '1'.
	zeros := 0
	for zeros < len(data) && data[zeros] == 0 {
		zeros++
	}
	// Big-endian byte array → big.Int (skipping the leading zeros).
	num := new(big.Int)
	for i := zeros; i < len(data); i++ {
		num.Lsh(num, 8)
		num.Or(num, big.NewInt(int64(data[i])))
	}
	// Base-convert to 58 digits (collected least-significant first).
	var digits []byte
	base := big.NewInt(58)
	v := new(big.Int).Set(num)
	for v.Sign() > 0 {
		mod := new(big.Int)
		v.DivMod(v, base, mod)
		digits = append(digits, byte(mod.Int64()))
	}
	var sb strings.Builder
	for i := 0; i < zeros; i++ {
		sb.WriteByte('1')
	}
	for i := len(digits) - 1; i >= 0; i-- {
		sb.WriteByte(b58Alphabet[digits[i]])
	}
	return sb.String()
}

func decode58(s string) ([]byte, error) {
	// Count leading '1's — each maps to a 0x00 byte.
	zeros := 0
	for zeros < len(s) && s[zeros] == '1' {
		zeros++
	}
	num := new(big.Int)
	base := big.NewInt(58)
	for i := zeros; i < len(s); i++ {
		idx := strings.IndexByte(b58Alphabet, s[i])
		if idx == -1 {
			return nil, errInvalid
		}
		num.Mul(num, base)
		num.Add(num, big.NewInt(int64(idx)))
	}
	// big.Int → big-endian bytes (minimal representation).
	body := num.Bytes()
	out := make([]byte, zeros+len(body))
	copy(out[zeros:], body)
	return out, nil
}

// ---------------------------------------------------------------------------
// Base62 — standard big.Int base-conversion of the byte array (no leading-zero
// special-casing beyond the standard big-int).
// ---------------------------------------------------------------------------

func encode62(data []byte) string {
	if len(data) == 0 {
		return ""
	}
	num := new(big.Int)
	for _, b := range data {
		num.Lsh(num, 8)
		num.Or(num, big.NewInt(int64(b)))
	}
	if num.Sign() == 0 {
		return "0"
	}
	var digits []byte
	base := big.NewInt(62)
	v := new(big.Int).Set(num)
	for v.Sign() > 0 {
		mod := new(big.Int)
		v.DivMod(v, base, mod)
		digits = append(digits, byte(mod.Int64()))
	}
	var sb strings.Builder
	for i := len(digits) - 1; i >= 0; i-- {
		sb.WriteByte(b62Alphabet[digits[i]])
	}
	return sb.String()
}

func decode62(s string) ([]byte, error) {
	if len(s) == 0 {
		return []byte{}, nil
	}
	num := new(big.Int)
	base := big.NewInt(62)
	for i := 0; i < len(s); i++ {
		idx := strings.IndexByte(b62Alphabet, s[i])
		if idx == -1 {
			return nil, errInvalid
		}
		num.Mul(num, base)
		num.Add(num, big.NewInt(int64(idx)))
	}
	return num.Bytes(), nil
}

// ---------------------------------------------------------------------------
// Base85 — Ascii85. 4 bytes → 5 chars in '!'(33)..'u'(117); a full 4-zero
// group is shortened to 'z'. No <~ ~> delimiters. Partial final groups emit
// one fewer char than (bytes+1) would suggest; decode reverses, padding with
// 'u' (value 84).
// ---------------------------------------------------------------------------

func encode85(data []byte) string {
	var sb strings.Builder
	for i := 0; i < len(data); i += 4 {
		end := i + 4
		if end > len(data) {
			end = len(data)
		}
		chunk := data[i:end]
		isFull := len(chunk) == 4
		var b [4]int
		for j := 0; j < len(chunk); j++ {
			b[j] = int(chunk[j])
		}
		u := b[0]*16777216 + b[1]*65536 + b[2]*256 + b[3]
		if isFull && u == 0 {
			sb.WriteByte('z') // zero-group shorthand
			continue
		}
		var digits [5]int
		v := u
		for k := 4; k >= 0; k-- {
			digits[k] = v % 85
			v /= 85
		}
		if isFull {
			for _, d := range digits {
				sb.WriteByte(byte(d + 33))
			}
		} else {
			n := len(chunk) + 1 // n bytes → n+1 chars
			for k := 0; k < n; k++ {
				sb.WriteByte(byte(digits[k] + 33))
			}
		}
	}
	return sb.String()
}

func decode85(s string) ([]byte, error) {
	var out []byte
	group := make([]int, 0, 5)
	for i := 0; i < len(s); i++ {
		c := s[i]
		if c == 'z' {
			// 'z' is only valid at a group boundary (an empty accumulator).
			if len(group) != 0 {
				return nil, errInvalid
			}
			out = append(out, 0, 0, 0, 0)
			continue
		}
		if c < 33 || c > 117 {
			return nil, errInvalid
		}
		group = append(group, int(c-33))
		if len(group) == 5 {
			v := 0
			for _, d := range group {
				v = v*85 + d
			}
			if v > 0xffffffff {
				return nil, errInvalid // a 5-char group must fit in 32 bits
			}
			out = append(out, byte((v>>24)&0xff), byte((v>>16)&0xff), byte((v>>8)&0xff), byte(v&0xff))
			group = group[:0]
		}
	}
	// Handle a partial final group (2–4 chars → 1–3 bytes).
	if len(group) > 0 {
		m := len(group)
		if m < 2 {
			return nil, errInvalid // a lone trailing char is malformed
		}
		for len(group) < 5 {
			group = append(group, 84) // pad with 'u'
		}
		v := 0
		for _, d := range group {
			v = v*85 + d
		}
		if v > 0xffffffff {
			return nil, errInvalid
		}
		all := [4]byte{byte((v >> 24) & 0xff), byte((v >> 16) & 0xff), byte((v >> 8) & 0xff), byte(v & 0xff)}
		out = append(out, all[:m-1]...)
	}
	return out, nil
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

// encodeBytes dispatches raw bytes to the chosen scheme's encoder.
func encodeBytes(data []byte, scheme Scheme) string {
	switch scheme {
	case Base32:
		return encode32(data)
	case Base58:
		return encode58(data)
	case Base62:
		return encode62(data)
	case Base85:
		return encode85(data)
	default:
		return ""
	}
}

// decodeBytes dispatches an encoded string to the chosen scheme's decoder. An
// invalid or malformed input yields errInvalid (mirroring the TS `null`).
func decodeBytes(encoded string, scheme Scheme) ([]byte, error) {
	switch scheme {
	case Base32:
		return decode32(encoded)
	case Base58:
		return decode58(encoded)
	case Base62:
		return decode62(encoded)
	case Base85:
		return decode85(encoded)
	default:
		return nil, errInvalid
	}
}

// Encode returns the chosen-scheme encoding of the UTF-8 bytes of text. Empty
// text encodes to "". It is the Go twin of encode() in src/lib/base-encoder.ts.
func Encode(text string, scheme Scheme) string {
	return encodeBytes([]byte(text), scheme)
}

// Decode reverses an encoded string back to UTF-8 text. Invalid characters or
// a malformed structure yield errInvalid — mirroring the TS lib's `null`. It
// is the Go twin of decode() in src/lib/base-encoder.ts.
//
// The decoded bytes are interpreted as UTF-8; for all shared round-trip
// vectors (which start from valid UTF-8 text) this matches the TS TextDecoder.
func Decode(encoded string, scheme Scheme) (string, error) {
	data, err := decodeBytes(encoded, scheme)
	if err != nil {
		return "", err
	}
	return string(data), nil
}