unigram 0.1.5

//! `unigram` — a bijective codec between bytes and words that cost one LLM token.
//!
//! Machine identifiers spend their lives being looked at: handed to a language model
//! and asked back, printed in a log, quoted in an error, read off a page by whoever
//! is debugging at the time. This crate carries them as words, so that an id becomes
//! something a reader can hold — `department number access world` can be said out
//! loud, told apart from its neighbour at a glance, and recognised again an hour
//! later, which is what a name is for.
//!
//! ```
//! let words = unigram::encode(&[0x3d, 0x9a, 0x00, 0xff]);
//! assert_eq!(words, "department number access world");
//! assert_eq!(unigram::decode(&words).unwrap(), vec![0x3d, 0x9a, 0x00, 0xff]);
//! ```
//!
//! It is also the densest form the trip allows. The words come from a fixed alphabet
//! of 256, and two properties follow from that size — they are the whole design:
//!
//! - **One word is exactly one byte.** Encoding is a table lookup per byte with no
//!   bit-packing, no padding, and no length convention; decoding is its inverse.
//!   Every byte string has exactly one encoding, and every sequence of alphabet
//!   words decodes.
//! - **Every word is exactly one token.** Each entry was measured to cost a single
//!   token, so an encoded value costs exactly one token per byte — and, unlike hex,
//!   the same for every value. Measured against hex of the same payload:
//!
//! | payload  | bits | hex (mean / worst) | `unigram` |
//! |----------|------|--------------------|-----------|
//! | 4 bytes  | 32   | 6.0 / 8            | 4         |
//! | 8 bytes  | 64   | 11.1 / 14          | 8         |
//! | 16 bytes | 128  | 21.5 / 25          | 16        |
//! | 32 bytes | 256  | 42.2 / 49          | 32        |
//!
//!   Roughly a quarter cheaper on average, but the flat cost matters more than the
//!   mean: hex cost swings with the value, so a budget built on it has to assume
//!   the worst case. Those are Claude's numbers; the margin narrows under the GPT-4
//!   vocabularies, where 32 bytes of hex average 37.1, and widens sharply under
//!   Llama's SentencePiece, where it averages 58.2 against the same flat 32.
//!   `verify-alphabet.py` prints the table for every family it checks.
//!
//!
//! ## Why the join is a space
//!
//! Tokenizer vocabularies hold their canonical word entries space-prefixed, so the
//! space between two words is absorbed into the word that follows it and costs
//! nothing. No other separator is free. Measured across all five families, a hyphen,
//! comma, pipe, slash, or newline becomes a token of its own in every one of them,
//! taking an eight-byte value from 8 tokens to 15 — the join costing almost as much
//! as the payload. GPT-3.5/4 and GPT-4o absorb `_` and `.` for free; no other family
//! absorbs anything. Encoded values travel inside quoted strings in practice, where
//! embedded spaces are free.
//!
//! [`decode`] is nonetheless liberal in what it accepts: any run of characters that
//! is not an ASCII letter separates words, and case is ignored. A value that came
//! back hyphenated, re-wrapped across lines, comma-joined, or shouted still decodes
//! to the bytes that were sent.
//!
//! ## Choosing a length
//!
//! One word is one byte and one token, so a value's length is its entropy budget
//! and its token budget at once — the two cannot drift apart, which is most of why
//! this is easier to reason about than hex.
//!
//! | words | bits | distinct values | values before a 1-in-a-million collision |
//! |-------|------|-----------------|------------------------------------------|
//! | 2     | 16   | 65,536          | fewer than 1                             |
//! | 3     | 24   | 16.8 million    | 5                                        |
//! | 4     | 32   | 4.3 billion     | 92                                       |
//! | 6     | 48   | 281 trillion    | 23,700                                   |
//! | 8     | 64   | 1.8 × 10^19     | 6 million                                |
//! | 16    | 128  | 3.4 × 10^38     | 2.6 × 10^16                              |
//! | 32    | 256  | 1.2 × 10^77     | 4.8 × 10^35                              |
//!
//! The right column is the birthday bound, `k ≈ sqrt(2·N·p)`, and it is the column
//! to size against: collisions arrive at the square root of the space, not at the
//! space. Sixteen words is a UUID's width, thirty-two a SHA-256's.
//!
//! Two questions hide in that table and it answers only one. **Collision** is the
//! right column — how many values may be outstanding before two coincide.
//! **Guessing** is separate: [`mint`] draws from the OS CSPRNG, so every bit is
//! unpredictable, but four words is 4.3 billion candidates, which is an afternoon
//! for anything that can ask freely. Four words suits a value that is scoped,
//! short-lived, and rate-limited — an acknowledgement nonce, a correlation id. A
//! value a stranger can grind at wants eight or more, and at equal entropy the words
//! are still the cheaper carrier — see the table above.
//!
//! ## The alphabet
//!
//! Entries are lowercase ASCII English, 4 to 11 characters, chosen under four
//! constraints:
//!
//! - **One token, under five tokenizer families.** Every entry costs a single token
//!   under Claude, GPT-2/3 (`r50k`, `p50k`), GPT-3.5/4 (`cl100k`), GPT-4o
//!   (`o200k`), and Llama's SentencePiece — spanning both the BPE and SentencePiece
//!   families. None of those vocabularies is vendored here, so this is checked by a
//!   script rather than by `cargo test`; see "Changing the alphabet" below.
//! - **No two entries within one character edit of each other.** A slipped character
//!   therefore lands outside the alphabet rather than on a different valid word, so
//!   [`decode`] refuses it instead of returning different bytes. This matters less
//!   than it sounds — neither a model nor a copy-paste mangles an id in practice —
//!   but it is free, given that the entries have to be distinguishable to read.
//! - **Nothing charged** — no death, violence, race, gender, religion, or politics.
//!   These strings surface unbidden in transcripts, logs, and user-facing errors.
//! - **No entry is an inflection of another**, so a dropped plural cannot silently
//!   decode to a different byte.
//!
//! ## Why the alphabet is 256 and not larger
//!
//! A wider alphabet would carry more bits per token, so it is worth saying why this
//! one stops where it does. Of the roughly 65,000 space-prefixed lowercase words in
//! the largest vocabulary, 6,654 are single-token in all five families; 5,452 of
//! those are 4 to 11 ASCII characters; and 640 of *those* survive Claude, whose
//! tokenizer is by far the narrowest of the five. Spacing them a character edit
//! apart leaves about 509.
//!
//! So the ceiling is 512 entries — 9 bits per token against the 8 here, and 9 does
//! not divide 8. Bit-packing 9-bit symbols would save nothing at all on a 4-byte
//! value (32 bits still needs 4 words), one token on a 16-byte value, and three on a
//! 32-byte one, in exchange for the byte-indexed table, the claim that one word is
//! one byte, and a codec that can be described in a sentence. It is not a trade
//! worth making.
//!
//! Other scripts do not change this. CJK is denser on the page but agrees across
//! families far less: 39 characters are single-token in all five, which does not
//! reach even 256. Accented Latin is worse — 5 words survive. The binding constraint
//! was never English; it is the intersection itself.
//!
//! ## Changing the alphabet
//!
//! Nothing here tokenizes, at runtime or under test: the OS CSPRNG is this crate's
//! only dependency at any stage. So `cargo test` covers the codec's behaviour and
//! the table's structural properties — 256 entries, sorted, unique, 4 to 11
//! lowercase ASCII characters, and no two within one edit — and says nothing about
//! cost.
//!
//! Every cost claim above is checked instead by `verify-alphabet.py`, beside this
//! file. It reads [`ALPHABET`] straight out of this source — a copy would drift —
//! and re-measures each entry against all five tokenizer families, along with the
//! composed per-byte cost, the margin over hex, and the choice of separator:
//!
//! ```text
//! uv run verify-alphabet.py
//! ```
//!
//! Run it after any edit to [`ALPHABET`]. A green test suite alone establishes none
//! of what this crate is named for, and an edit that satisfies every test here can
//! still break every cost claim above.

#![forbid(unsafe_code)]

use std::fmt;

/// The 256-word alphabet, sorted, indexed by the byte each word encodes.
///
/// Sorted so [`decode`] can binary-search it, and byte `n` is `ALPHABET[n]` — the
/// table *is* the codec. Reordering an entry changes what every previously issued
/// value decodes to, so this list is appended to, never rearranged.
///
/// Laid out packed rather than one entry per line: rustfmt would give this table
/// 256 vertical lines, which is harder to scan and to review than a grid, and the
/// entries are data rather than code.
///
/// Editing this table? A green `cargo test` proves only its structure. Run
/// `verify-alphabet.py` — that is where single-token cost is checked, under all
/// five tokenizer families.
#[rustfmt::skip]
pub const ALPHABET: [&str; 256] = [
    "access", "account", "action", "address", "album", "android", "application", "area",
    "array", "article", "association", "author", "award", "background", "band", "black",
    "board", "body", "border", "break", "build", "building", "business", "button", "call",
    "card", "career", "category", "census", "center", "central", "century", "change", "character",
    "check", "city", "class", "click", "client", "close", "club", "code", "college", "color",
    "column", "command", "common", "community", "company", "components", "console", "container",
    "content", "control", "council", "count", "country", "course", "data", "database", "density",
    "department", "description", "design", "development", "device", "director", "display",
    "district", "division", "document", "door", "double", "download", "early", "education",
    "element", "email", "error", "events", "example", "export", "express", "face", "features",
    "field", "film", "first", "float", "food", "football", "force", "form", "format", "function",
    "future", "games", "general", "green", "group", "head", "header", "height", "help", "high",
    "history", "home", "host", "house", "households", "images", "import", "important", "income",
    "index", "info", "information", "input", "install", "island", "king", "label", "language",
    "large", "league", "length", "level", "library", "license", "life", "light", "list",
    "local", "location", "login", "love", "management", "march", "market", "master", "material",
    "math", "median", "members", "message", "method", "million", "models", "money", "music",
    "network", "news", "north", "note", "number", "object", "office", "options", "package",
    "page", "password", "people", "period", "person", "places", "play", "players", "population",
    "port", "position", "power", "press", "price", "print", "println", "process", "production",
    "products", "program", "project", "property", "published", "query", "question", "range",
    "records", "references", "region", "register", "render", "report", "request", "research",
    "response", "results", "return", "review", "river", "role", "room", "router", "school",
    "science", "score", "script", "search", "season", "section", "select", "send", "series",
    "services", "session", "share", "social", "society", "software", "song", "source", "south",
    "space", "span", "species", "square", "station", "story", "street", "string", "students",
    "study", "style", "success", "system", "table", "target", "task", "team", "television",
    "template", "title", "token", "track", "train", "training", "union", "university", "update",
    "username", "users", "version", "video", "village", "website", "width", "window", "world",
];

/// Why a sequence of words could not be decoded.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecodeError {
    /// A word outside the alphabet, and where in the sequence it sat.
    ///
    /// Reported rather than skipped or guessed at: a value that lost a word is not
    /// the value that was sent, and inventing the byte it stood for would answer a
    /// question nobody asked with a value nobody issued.
    UnknownWord { position: usize, word: String },
    /// No alphabet words at all. The encoding of no bytes is the empty string, which
    /// is never a value a caller means to transmit, so decoding one is an error
    /// rather than an empty success.
    Empty,
}

impl fmt::Display for DecodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnknownWord { position, word } => write!(
                f,
                "`{word}` (word {}) is not in the unigram alphabet",
                position + 1
            ),
            Self::Empty => f.write_str("no unigram words found"),
        }
    }
}

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

/// Split on any run of characters that is not an ASCII letter.
///
/// Being this liberal is what lets a value survive a round trip through a model or
/// a transcript: hyphens, newlines, commas, quotes, and stray punctuation all read
/// as separators, so only the words themselves have to arrive intact.
fn split_words(text: &str) -> impl Iterator<Item = &str> {
    text.split(|c: char| !c.is_ascii_alphabetic())
        .filter(|word| !word.is_empty())
}

/// Encode bytes as space-joined alphabet words, one word per byte.
///
/// Cheap enough to call at a boundary rather than storing the result: four bytes
/// become twenty-six characters, which is a poor thing to keep in a column, and this
/// is a table lookup per byte in each direction. Store the bytes; render the words
/// wherever they will actually be read.
pub fn encode(bytes: &[u8]) -> String {
    let mut out = String::with_capacity(bytes.len() * 8);
    for (index, byte) in bytes.iter().enumerate() {
        if index > 0 {
            out.push(' ');
        }
        out.push_str(ALPHABET[*byte as usize]);
    }
    out
}

/// Decode alphabet words back to the bytes they carry.
///
/// Liberal in what it accepts — any run of characters that is not an ASCII letter
/// separates words, and case is ignored — but exact in what it returns: every word
/// must be in the alphabet, or the value is refused and the offending word named.
pub fn decode(text: &str) -> Result<Vec<u8>, DecodeError> {
    let mut bytes = Vec::new();
    for (position, word) in split_words(text).enumerate() {
        let lowered = word.to_ascii_lowercase();
        match ALPHABET.binary_search(&lowered.as_str()) {
            Ok(index) => bytes.push(index as u8),
            Err(_) => {
                return Err(DecodeError::UnknownWord {
                    position,
                    word: word.to_string(),
                })
            }
        }
    }
    if bytes.is_empty() {
        return Err(DecodeError::Empty);
    }
    Ok(bytes)
}

/// Mint `bytes` bytes of fresh entropy, encoded.
///
/// Four bytes is a reasonable default for a short-lived, rate-limited nonce: 32
/// bits in four flat tokens, where the same 32 bits as hex average 6 and can reach
/// 8. It is a poor default for anything else — see "Choosing a length" above, which
/// is the difference between a value that cannot collide and one that cannot be
/// guessed.
///
/// # Panics
///
/// If the OS entropy source is unavailable. That is not a condition a caller can
/// do anything useful with, and returning a predictable value instead would be far
/// worse than stopping.
pub fn mint(bytes: usize) -> String {
    let mut buffer = vec![0u8; bytes];
    getrandom::fill(&mut buffer).expect("OS entropy source unavailable");
    encode(&buffer)
}

/// Reduce a value to the form comparisons are made in: trimmed, lowercased, and
/// with internal whitespace runs collapsed to a single space.
///
/// Deliberately preserves every non-whitespace character, so this is safe to apply
/// to a string that is *not* an encoded value — a legacy hex token, say — without
/// mangling it. [`decode`]'s liberal splitting is the opposite trade and belongs
/// only where the bytes are actually wanted back.
pub fn normalize(text: &str) -> String {
    text.split_whitespace()
        .map(|part| part.to_ascii_lowercase())
        .collect::<Vec<_>>()
        .join(" ")
}

/// Compare a value that was issued against one a caller presented, tolerating the
/// damage a round trip through a model or a transcript does.
///
/// When both sides are encoded values the comparison is on the decoded bytes, so
/// separator and case damage on the presented side cannot matter. Otherwise it
/// falls back to comparing [`normalize`]d strings, which is what lets a value
/// issued in some older format still match itself without a migration.
pub fn matches(issued: &str, presented: &str) -> bool {
    if let (Ok(issued_bytes), Ok(presented_bytes)) = (decode(issued), decode(presented)) {
        return issued_bytes == presented_bytes;
    }
    normalize(issued) == normalize(presented)
}

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

    /// Everything cost-related lives in `verify-alphabet.py`, which measures the
    /// five tokenizer families this crate claims. What is left here is what can be
    /// checked without a vocabulary: the table's structure, and the codec over it.
    #[test]
    fn the_alphabet_is_sorted_unique_and_plain_lowercase() {
        let mut sorted = ALPHABET;
        sorted.sort_unstable();
        assert_eq!(sorted, ALPHABET, "binary_search requires sorted order");
        let unique: std::collections::HashSet<_> = ALPHABET.iter().collect();
        assert_eq!(unique.len(), ALPHABET.len());
        for word in ALPHABET {
            assert!(
                word.len() >= 4 && word.len() <= 11 && word.bytes().all(|b| b.is_ascii_lowercase()),
                "`{word}`"
            );
        }
    }

    /// Distance from every other entry is what turns a one-character slip into a
    /// refusal instead of a different valid byte. Without it the codec is only as
    /// honest as hex.
    #[test]
    fn no_two_entries_are_within_one_edit_of_each_other() {
        fn within_one_edit(a: &str, b: &str) -> bool {
            let (a, b) = if a.len() > b.len() { (b, a) } else { (a, b) };
            let (short, long) = (a.as_bytes(), b.as_bytes());
            match long.len() - short.len() {
                0 => short.iter().zip(long).filter(|(x, y)| x != y).count() <= 1,
                1 => {
                    let skip = short.iter().zip(long).take_while(|(x, y)| x == y).count();
                    short[skip..] == long[skip + 1..]
                }
                _ => false,
            }
        }
        for (i, a) in ALPHABET.iter().enumerate() {
            for b in &ALPHABET[i + 1..] {
                assert!(!within_one_edit(a, b), "`{a}` and `{b}` are one edit apart");
            }
        }
    }

    #[test]
    fn every_byte_round_trips() {
        let all: Vec<u8> = (0..=255).collect();
        assert_eq!(decode(&encode(&all)).unwrap(), all);
    }

    #[test]
    fn a_single_byte_round_trips_without_separators() {
        let encoded = encode(&[7]);
        assert!(!encoded.contains(' '));
        assert_eq!(decode(&encoded).unwrap(), vec![7]);
    }

    /// The point of the codec: a value mangled on its way through a model still
    /// decodes to what was sent.
    #[test]
    fn decoding_survives_the_mangling_a_round_trip_introduces() {
        let bytes = [0x3d, 0x9a, 0x00, 0xff];
        let encoded = encode(&bytes);
        for mangled in [
            encoded.to_uppercase(),
            format!("  {encoded}  "),
            encoded.replace(' ', "-"),
            encoded.replace(' ', ",  "),
            encoded.replace(' ', "\n"),
            format!("\"{}\"", encoded.replace(' ', "   ")),
        ] {
            assert_eq!(decode(&mangled).unwrap(), bytes, "{mangled}");
        }
    }

    #[test]
    fn an_unknown_word_is_refused_and_named() {
        let encoded = format!("{} zzzz {}", ALPHABET[1], ALPHABET[2]);
        assert_eq!(
            decode(&encoded),
            Err(DecodeError::UnknownWord {
                position: 1,
                word: "zzzz".to_string(),
            })
        );
    }

    /// A near-miss is the case that matters: one character off a real entry must be
    /// refused, not silently read as some other byte.
    #[test]
    fn a_one_character_slip_is_refused_rather_than_read_as_another_byte() {
        assert!(matches!(
            decode("accesx"),
            Err(DecodeError::UnknownWord { .. })
        ));
    }

    #[test]
    fn an_empty_value_is_refused() {
        assert_eq!(decode(""), Err(DecodeError::Empty));
        assert_eq!(decode("   -- \n"), Err(DecodeError::Empty));
    }

    #[test]
    fn mint_produces_one_word_per_requested_byte() {
        let minted = mint(4);
        assert_eq!(minted.split(' ').count(), 4, "{minted}");
        assert_eq!(decode(&minted).unwrap().len(), 4);
        assert_ne!(mint(8), mint(8));
    }

    #[test]
    fn matching_tolerates_mangling_of_an_encoded_value() {
        let issued = mint(4);
        assert!(matches(&issued, &issued));
        assert!(matches(&issued, &issued.to_uppercase()));
        assert!(matches(
            &issued,
            &format!("  {}  ", issued.replace(' ', " - "))
        ));
        assert!(!matches(&issued, &mint(4)));
    }

    /// Values issued in an older opaque format have to keep matching themselves, or
    /// swapping the minted form would strand every token outstanding at the moment
    /// of the upgrade.
    #[test]
    fn matching_still_compares_values_that_are_not_encoded_at_all() {
        let legacy = "3925ca9a0065442496cc231d6ae48870";
        assert!(matches(legacy, legacy));
        assert!(matches(legacy, &format!("  {}  ", legacy.to_uppercase())));
        assert!(!matches(legacy, "3925ca9a0065442496cc231d6ae48871"));
        assert!(!matches(legacy, &mint(4)));
    }

    #[test]
    fn normalize_leaves_a_non_encoded_string_intact() {
        assert_eq!(normalize("  3925CA9A-0065  "), "3925ca9a-0065");
    }
}