tokenx 0.2.0

//! Per-segment token counting and the option types.

use crate::patterns;
use regex::{Regex, RegexBuilder};
use std::sync::OnceLock;

/// A language rule. When `pattern` matches a segment, the segment uses
/// `average_chars_per_token` instead of the default.
#[derive(Clone, Debug)]
pub struct LanguageConfig {
    /// Regex that detects the language. It fires if it matches anywhere in the
    /// segment.
    pub pattern: Regex,
    /// Average characters per token for segments this rule matches.
    pub average_chars_per_token: f64,
}

impl LanguageConfig {
    /// Builds a config from a pattern source and a chars-per-token value.
    ///
    /// The pattern compiles case sensitive. Use [`LanguageConfig::case_insensitive`]
    /// to fold case the way a JavaScript `/i` regex does.
    ///
    /// # Errors
    ///
    /// Returns an error if `pattern` is not a valid regex.
    pub fn new(pattern: &str, average_chars_per_token: f64) -> Result<Self, regex::Error> {
        Ok(Self {
            pattern: Regex::new(pattern)?,
            average_chars_per_token,
        })
    }

    /// Builds a config whose pattern folds case, matching a `/i` JavaScript regex.
    ///
    /// # Errors
    ///
    /// Returns an error if `pattern` is not a valid regex.
    pub fn case_insensitive(
        pattern: &str,
        average_chars_per_token: f64,
    ) -> Result<Self, regex::Error> {
        Ok(Self {
            pattern: RegexBuilder::new(pattern).case_insensitive(true).build()?,
            average_chars_per_token,
        })
    }
}

/// Options for token estimation.
///
/// A `None` field uses the built-in default. A `Some` value is used as given,
/// including `Some(0.0)`. This matches JavaScript destructuring defaults, which
/// apply only when a property is absent.
#[derive(Clone, Debug, Default)]
pub struct TokenEstimationOptions {
    /// Default average characters per token. `None` uses
    /// [`patterns::DEFAULT_CHARS_PER_TOKEN`].
    pub default_chars_per_token: Option<f64>,
    /// Language rules. `None` uses the built-in set of three.
    pub language_configs: Option<Vec<LanguageConfig>>,
}

/// Options for [`crate::split_by_tokens`]. Adds chunk overlap.
#[derive(Clone, Debug, Default)]
pub struct SplitByTokensOptions {
    /// Default average characters per token. `None` uses
    /// [`patterns::DEFAULT_CHARS_PER_TOKEN`].
    pub default_chars_per_token: Option<f64>,
    /// Language rules. `None` uses the built-in set of three.
    pub language_configs: Option<Vec<LanguageConfig>>,
    /// Tokens to repeat between consecutive chunks. `None` means zero.
    pub overlap: Option<u64>,
}

/// The three built-in language rules, compiled once and shared.
///
/// Order matters. The first rule whose pattern matches wins. German and the
/// Romance set use 3 chars per token. The Slavic set uses 3.5. The set is
/// cached so default-options calls do not recompile the regexes each time.
pub fn default_language_configs() -> &'static [LanguageConfig] {
    static CONFIGS: OnceLock<Vec<LanguageConfig>> = OnceLock::new();
    CONFIGS
        .get_or_init(|| {
            vec![
                LanguageConfig::case_insensitive("[äöüßẞ]", 3.0).unwrap(),
                LanguageConfig::case_insensitive("[éèêëàâîïôûùüÿçœæáíóúñ]", 3.0).unwrap(),
                LanguageConfig::case_insensitive("[ąćęłńóśźżěščřžýůúďťň]", 3.5).unwrap(),
            ]
        })
        .as_slice()
}

/// UTF-16 code-unit length, the unit JavaScript `String.length` reports.
pub fn utf16_len(s: &str) -> usize {
    s.encode_utf16().count()
}

/// Unicode code-point count, the unit `Array.from(s).length` reports.
pub fn char_count(s: &str) -> usize {
    s.chars().count()
}

/// Counts tokens in one segment.
///
/// Rules apply in order and the first match wins:
/// 1. whitespace -> 0
/// 2. contains CJK -> one token per code point
/// 3. whole-segment numeric -> 1
/// 4. UTF-16 length <= 3 -> 1
/// 5. contains punctuation -> ceil(len / 2)
/// 6. whole-segment alphanumeric -> `ceil(len / chars_per_token)`
/// 7. fallback (mixed content) -> `ceil(len / chars_per_token)`
pub fn estimate_segment_tokens(
    segment: &str,
    language_configs: &[LanguageConfig],
    default_chars_per_token: f64,
) -> u64 {
    if patterns::whitespace().is_match(segment) {
        return 0;
    }

    if patterns::cjk().is_match(segment) {
        return char_count(segment) as u64;
    }

    if patterns::numeric().is_match(segment) {
        return 1;
    }

    let len = utf16_len(segment);
    if len <= patterns::SHORT_TOKEN_THRESHOLD {
        return 1;
    }

    if patterns::punctuation().is_match(segment) {
        return (len as f64 / 2.0).ceil() as u64;
    }

    // Rules 6 and 7 (whole-segment alphanumeric, mixed-content fallback) compute
    // the same value, so the alphanumeric test does not change the result and is
    // not run.
    let chars_per_token = language_specific_chars_per_token(segment, language_configs)
        .unwrap_or(default_chars_per_token);
    // A zero or negative chars-per-token makes the division infinite or negative.
    // `as u64` saturates: infinity becomes u64::MAX, a negative becomes 0. The
    // result stays finite and no arithmetic panics.
    (len as f64 / chars_per_token).ceil() as u64
}

/// Returns the chars-per-token of the first matching language rule, or `None`.
fn language_specific_chars_per_token(
    segment: &str,
    language_configs: &[LanguageConfig],
) -> Option<f64> {
    language_configs
        .iter()
        .find(|c| c.pattern.is_match(segment))
        .map(|c| c.average_chars_per_token)
}

/// Splits text into segments, keeping whitespace and punctuation runs as their
/// own segments and dropping empty pieces. A dot or comma inside a standalone
/// ASCII-digit number stays in that numeric segment.
///
/// This matches JavaScript `String.split` with a single capturing group:
/// delimiters stay in the output, and empty strings between adjacent delimiters
/// are removed, except for the numeric separator rule above.
pub fn split_segments(text: &str) -> Vec<&str> {
    let mut out = Vec::new();
    let mut start = 0;
    let mut i = 0;

    while i < text.len() {
        if i == start
            && text.as_bytes()[i].is_ascii_digit()
            && scan_boundary_number(text, &mut start, &mut i, &mut out)
        {
            continue;
        }

        if is_delimiter_at(text, i) {
            if i > start {
                out.push(&text[start..i]);
            }
            let delimiter_start = i;
            let punctuation = is_punctuation(text.as_bytes()[i]);
            i += delimiter_len_at(text, i);
            while i < text.len()
                && ((punctuation && is_punctuation(text.as_bytes()[i]))
                    || (!punctuation && is_whitespace_at(text, i)))
            {
                i += delimiter_len_at(text, i);
            }
            out.push(&text[delimiter_start..i]);
            start = i;
        } else {
            i += char_len_at(text, i);
        }
    }

    if start < text.len() {
        out.push(&text[start..]);
    }

    out
}

fn scan_boundary_number<'a>(
    text: &'a str,
    start: &mut usize,
    i: &mut usize,
    out: &mut Vec<&'a str>,
) -> bool {
    let original_start = *start;
    let original_len = out.len();
    let bytes = text.as_bytes();

    loop {
        let digit_start = *i;
        consume_ascii_digits(bytes, i);

        if *i < bytes.len()
            && matches!(bytes[*i], b'.' | b',')
            && *i + 1 < bytes.len()
            && bytes[*i + 1].is_ascii_digit()
        {
            out.push(&text[digit_start..*i]);
            out.push(&text[*i..*i + 1]);
            *i += 1;
            continue;
        }

        let standalone_end = *i == text.len() || is_delimiter_at(text, *i);
        if standalone_end && patterns::numeric().is_match(&text[original_start..*i]) {
            out.truncate(original_len);
            out.push(&text[original_start..*i]);
            *start = *i;
        } else if out.len() == original_len {
            *start = original_start;
        } else {
            *start = digit_start;
        }

        return true;
    }
}

fn consume_ascii_digits(bytes: &[u8], i: &mut usize) {
    while *i < bytes.len() && bytes[*i].is_ascii_digit() {
        *i += 1;
    }
}

fn is_delimiter_at(text: &str, i: usize) -> bool {
    is_punctuation(text.as_bytes()[i]) || is_whitespace_at(text, i)
}

fn is_whitespace_at(text: &str, i: usize) -> bool {
    is_ecmascript_whitespace(text[i..].chars().next().unwrap())
}

fn delimiter_len_at(text: &str, i: usize) -> usize {
    if is_punctuation(text.as_bytes()[i]) {
        1
    } else {
        char_len_at(text, i)
    }
}

fn char_len_at(text: &str, i: usize) -> usize {
    text[i..].chars().next().unwrap().len_utf8()
}

fn is_punctuation(byte: u8) -> bool {
    matches!(
        byte,
        b'.' | b','
            | b'!'
            | b'?'
            | b';'
            | b'('
            | b')'
            | b'{'
            | b'}'
            | b'['
            | b']'
            | b'<'
            | b'>'
            | b':'
            | b'/'
            | b'\\'
            | b'|'
            | b'@'
            | b'#'
            | b'$'
            | b'%'
            | b'^'
            | b'&'
            | b'*'
            | b'+'
            | b'='
            | b'`'
            | b'~'
            | b'_'
            | b'-'
    )
}

fn is_ecmascript_whitespace(ch: char) -> bool {
    if ('\u{2000}'..='\u{200A}').contains(&ch) {
        return true;
    }

    matches!(
        ch,
        '\u{0009}'
            | '\u{000A}'
            | '\u{000B}'
            | '\u{000C}'
            | '\u{000D}'
            | '\u{0020}'
            | '\u{00A0}'
            | '\u{1680}'
            | '\u{2028}'
            | '\u{2029}'
            | '\u{202F}'
            | '\u{205F}'
            | '\u{3000}'
            | '\u{FEFF}'
    )
}

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

    fn est(segment: &str) -> u64 {
        estimate_segment_tokens(
            segment,
            default_language_configs(),
            patterns::DEFAULT_CHARS_PER_TOKEN,
        )
    }

    #[test]
    fn segment_math_table() {
        // Single-segment inputs and their direct token counts. Values match the
        // classifier rules in order.
        let cases: &[(&str, u64)] = &[
            ("is", 1),       // short token (UTF-16 length <= 3)
            ("a", 1),        // short token
            (",", 1),        // single punctuation, caught by short-token rule
            ("!!!", 1),      // length 3, short-token rule wins before punctuation
            ("!!!!", 2),     // length 4 punctuation, ceil(4 / 2)
            ("::::", 2),     // length 4 punctuation, ceil(4 / 2)
            ("12345", 1),    // numeric
            ("3.14", 1),     // numeric with a dot
            ("1,000.50", 1), // numeric with separators
            ("Hello", 1),    // alphanumeric, ceil(5 / 6)
            ("sentence", 2), // alphanumeric, ceil(8 / 6)
            ("道德經", 3),   // CJK, one token per code point
            // Astral characters inside a CJK segment. The CJK rule counts code
            // points, not UTF-16 units. Both inputs are 3 code points and 4
            // UTF-16 units, so they pin the code-point path. A switch to
            // utf16_len would return 4 and fail here.
            ("中😀文", 3),        // emoji between two Han characters
            ("中\u{20000}文", 3), // astral Han (Plane 2) between two BMP Han
            ("pünktlich", 3),     // German rule, ceil(9 / 3)
            ("français", 3),      // Romance rule, ceil(8 / 3)
            ("señor", 2),         // Romance rule, ceil(5 / 3)
            ("źdźbło", 2),        // Slavic rule, ceil(6 / 3.5)
            ("příliš", 2),        // Slavic rule, ceil(6 / 3.5)
        ];
        for (input, want) in cases {
            assert_eq!(est(input), *want, "segment {input:?}");
        }
    }

    #[test]
    fn whitespace_segment_is_zero() {
        assert_eq!(est("   "), 0);
        assert_eq!(est("\t\n "), 0);
    }

    #[test]
    fn uppercase_accents_match_language_rules() {
        // The default rules fold case, so uppercase accented letters match.
        // "ÄÖÜ" is length 3, so it hits the short-token rule and counts as 1.
        // A longer German word proves the rule fires.
        assert_eq!(est("ÜBERGROSS"), (9.0_f64 / 3.0).ceil() as u64);
    }

    #[test]
    fn split_keeps_delimiters() {
        let segs = split_segments("Hello, world!");
        assert_eq!(segs, vec!["Hello", ",", " ", "world", "!"]);
    }

    #[test]
    fn split_drops_empty_pieces() {
        // Adjacent delimiters and leading or trailing delimiters produce no
        // empty segments.
        let segs = split_segments("  a  ");
        assert_eq!(segs, vec!["  ", "a", "  "]);
    }
}