punycode-rs 0.1.0

//! A Rust port of Punycode (RFC 3492), the ASCII-compatible encoding at the heart of
//! Internationalized Domain Names — e.g. `münchen` <-> `mnchen-3ya`.
//!
//! Port target: Python's standard library `encodings/punycode.py` (Martin v. Löwis).
//! Verified against Python's built-in `punycode` codec.
//!
//! The algorithm separates a string into its ASCII "base" and its sorted non-ASCII
//! characters, then encodes each non-ASCII character as a single integer "delta" that
//! records both *which* character and *where* it is inserted. Deltas are written as
//! base-36 generalized variable-length integers with an adaptive bias that self-tunes
//! for compactness. Decoding mirrors the process exactly.

use std::fmt;

const BASE: i64 = 36;
const TMIN: i64 = 1;
const TMAX: i64 = 26;
const INITIAL_BIAS: i64 = 72;
const DAMP: i64 = 700;
const SKEW: i64 = 38;

/// `a-z0-9` — the 36 base-36 digit characters.
const DIGITS: &[u8; 36] = b"abcdefghijklmnopqrstuvwxyz0123456789";

/// An error from [`decode`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PunycodeError {
    /// The input contained a non-ASCII byte (Punycode is ASCII-only).
    NotAscii,
    /// The extended part ended mid-number.
    Incomplete,
    /// The extended part contained a byte that is not a valid base-36 digit.
    InvalidDigit,
    /// Decoding produced a code point that is not a valid character.
    InvalidCodePoint,
}

impl fmt::Display for PunycodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let msg = match self {
            PunycodeError::NotAscii => "input is not ASCII",
            PunycodeError::Incomplete => "incomplete punycode string",
            PunycodeError::InvalidDigit => "invalid extended code point",
            PunycodeError::InvalidCodePoint => "invalid decoded character",
        };
        f.write_str(msg)
    }
}

impl std::error::Error for PunycodeError {}

// --- Shared: threshold and bias adaptation. ---

/// `T(j, bias)` — the digit threshold, clamped to `[tmin, tmax]`.
fn threshold(j: i64, bias: i64) -> i64 {
    (BASE * (j + 1) - bias).clamp(TMIN, TMAX)
}

/// Bias adaptation after encoding/decoding a delta. Port of `adapt`.
fn adapt(mut delta: i64, first: bool, numchars: i64) -> i64 {
    delta /= if first { DAMP } else { 2 };
    delta += delta / numchars;
    let mut divisions = 0;
    while delta > ((BASE - TMIN) * TMAX) / 2 {
        delta /= BASE - TMIN;
        divisions += BASE;
    }
    divisions + (BASE * delta) / (delta + SKEW)
}

// --- Encoding (port of `punycode_encode`). ---

/// 3.1 Basic code point segregation: the ASCII bytes, and the sorted unique non-ASCII chars.
fn segregate(text: &[char]) -> (Vec<u8>, Vec<char>) {
    let mut base = Vec::new();
    let mut extended = std::collections::BTreeSet::new();
    for &c in text {
        if (c as u32) < 128 {
            base.push(c as u8);
        } else {
            extended.insert(c);
        }
    }
    (base, extended.into_iter().collect())
}

/// Count characters in `text` whose code point is below `max`.
fn selective_len(text: &[char], max: u32) -> i64 {
    text.iter().filter(|&&c| (c as u32) < max).count() as i64
}

/// Next occurrence of `target` in `text` after `pos`, returning `(index, pos)` where `index`
/// counts only characters at or below `target`. `(-1, -1)` when there is no further match.
fn selective_find(text: &[char], target: char, mut index: i64, mut pos: i64) -> (i64, i64) {
    let len = text.len() as i64;
    loop {
        pos += 1;
        if pos == len {
            return (-1, -1);
        }
        let c = text[pos as usize];
        if c == target {
            return (index + 1, pos);
        } else if c < target {
            index += 1;
        }
    }
}

/// 3.2 Insertion unsort coding: the sequence of deltas describing the non-ASCII insertions.
fn insertion_unsort(text: &[char], extended: &[char]) -> Vec<i64> {
    let mut oldchar: i64 = 0x80;
    let mut oldindex: i64 = -1;
    let mut result = Vec::new();

    for &c in extended {
        let mut index: i64 = -1;
        let mut pos: i64 = -1;
        let ch = c as i64;
        let curlen = selective_len(text, c as u32);
        let mut delta = (curlen + 1) * (ch - oldchar);
        loop {
            let (i, p) = selective_find(text, c, index, pos);
            index = i;
            pos = p;
            if index == -1 {
                break;
            }
            delta += index - oldindex;
            result.push(delta - 1);
            oldindex = index;
            delta = 0;
        }
        oldchar = ch;
    }
    result
}

/// 3.3 Generalized variable-length integer for one delta `n`.
fn generate_generalized_integer(mut n: i64, bias: i64) -> Vec<u8> {
    let mut result = Vec::new();
    let mut j = 0;
    loop {
        let t = threshold(j, bias);
        if n < t {
            result.push(DIGITS[n as usize]);
            return result;
        }
        result.push(DIGITS[(t + (n - t) % (BASE - t)) as usize]);
        n = (n - t) / (BASE - t);
        j += 1;
    }
}

/// 3.4 Encode all deltas, adapting the bias between each.
fn generate_integers(baselen: i64, deltas: &[i64]) -> Vec<u8> {
    let mut result = Vec::new();
    let mut bias = INITIAL_BIAS;
    for (points, &delta) in deltas.iter().enumerate() {
        result.extend(generate_generalized_integer(delta, bias));
        bias = adapt(delta, points == 0, baselen + points as i64 + 1);
    }
    result
}

/// Encode a string to Punycode (RFC 3492). Pure-ASCII input gains a trailing `-`.
pub fn encode(text: &str) -> String {
    let chars: Vec<char> = text.chars().collect();
    let (mut base, extended) = segregate(&chars);
    let deltas = insertion_unsort(&chars, &extended);
    let encoded = generate_integers(base.len() as i64, &deltas);

    if !base.is_empty() {
        base.push(b'-');
    }
    base.extend(encoded);
    String::from_utf8(base).expect("punycode output is ASCII")
}

// --- Decoding (port of `punycode_decode`). ---

/// 3.3 Decode one generalized integer, returning the new position and the value.
fn decode_generalized_number(
    extended: &[u8],
    mut extpos: usize,
    bias: i64,
) -> Result<(usize, i64), PunycodeError> {
    let mut result = 0;
    let mut w = 1;
    let mut j = 0;
    loop {
        let ch = *extended.get(extpos).ok_or(PunycodeError::Incomplete)?;
        extpos += 1;
        let digit = if ch.is_ascii_uppercase() {
            (ch - b'A') as i64
        } else if ch.is_ascii_digit() {
            (ch - b'0') as i64 + 26
        } else {
            return Err(PunycodeError::InvalidDigit);
        };
        let t = threshold(j, bias);
        result += digit * w;
        if digit < t {
            return Ok((extpos, result));
        }
        w *= BASE - t;
        j += 1;
    }
}

/// 3.2 Rebuild the string by decoding deltas and inserting characters. Port of `insertion_sort`.
fn insertion_sort(mut base: Vec<char>, extended: &[u8]) -> Result<Vec<char>, PunycodeError> {
    let mut char_code: i64 = 0x80;
    let mut pos: i64 = -1;
    let mut bias = INITIAL_BIAS;
    let mut extpos = 0;

    while extpos < extended.len() {
        let (newpos, delta) = decode_generalized_number(extended, extpos, bias)?;
        pos += delta + 1;
        char_code += pos / (base.len() as i64 + 1);
        if char_code > 0x10FFFF {
            return Err(PunycodeError::InvalidCodePoint);
        }
        pos %= base.len() as i64 + 1;
        let ch = char::from_u32(char_code as u32).ok_or(PunycodeError::InvalidCodePoint)?;
        base.insert(pos as usize, ch);
        bias = adapt(delta, extpos == 0, base.len() as i64);
        extpos = newpos;
    }
    Ok(base)
}

/// Decode a Punycode string back to Unicode. Input must be ASCII.
pub fn decode(text: &str) -> Result<String, PunycodeError> {
    if !text.is_ascii() {
        return Err(PunycodeError::NotAscii);
    }
    let bytes = text.as_bytes();

    // Split at the last '-': everything before is verbatim ASCII base, after is the
    // (case-insensitive) extended part.
    let (base, extended) = match bytes.iter().rposition(|&b| b == b'-') {
        Some(pos) => (
            bytes[..pos].iter().map(|&b| b as char).collect(),
            bytes[pos + 1..].to_ascii_uppercase(),
        ),
        None => (Vec::new(), bytes.to_ascii_uppercase()),
    };

    Ok(insertion_sort(base, &extended)?.into_iter().collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn encodes_known_values() {
        assert_eq!(encode("münchen"), "mnchen-3ya");
        assert_eq!(encode("abc"), "abc-"); // pure ASCII gains a trailing dash
        assert_eq!(encode(""), "");
        assert_eq!(encode("ü"), "tda");
    }

    #[test]
    fn decodes_known_values() {
        assert_eq!(decode("mnchen-3ya").unwrap(), "münchen");
        assert_eq!(decode("abc-").unwrap(), "abc");
        assert_eq!(decode("").unwrap(), "");
        assert_eq!(decode("tda").unwrap(), "ü");
    }

    #[test]
    fn round_trips() {
        for s in [
            "münchen",
            "café",
            "naïve",
            "日本語",
            "Ελληνικά",
            "abc",
            "Hello-World",
            "",
            "ñ",
            "a1b2c3",
            "emoji-💡-here",
            "mixed café 日本 test",
        ] {
            let encoded = encode(s);
            assert_eq!(decode(&encoded).unwrap(), s, "round-trip failed for {s:?}");
        }
    }

    #[test]
    fn decode_rejects_bad_input() {
        assert_eq!(decode("café"), Err(PunycodeError::NotAscii)); // non-ASCII input
        assert!(decode("-!").is_err()); // '!' is not a base-36 digit
    }
}