pos_tagger/

lexer.rs

//! Recursive regex tokenizer.
//!
//! The lexer splits a string against an ordered list of seven regexes. Each
//! level keeps the matched delimiters as their own tokens and re-splits the gaps
//! with the remaining regexes. Order sets precedence: URLs, IDs, and emails are
//! captured whole before numbers, whitespace, and single punctuation marks get a
//! chance to split.

use regex::Regex;
use std::sync::OnceLock;

/// Whitespace character class shared by every splitting regex.
///
/// This is the JavaScript whitespace set, the same code points
/// [`is_js_whitespace`] encodes. It differs from Rust's Unicode `\s` on two
/// code points. It excludes U+0085 (NEL) and includes U+FEFF (BOM). Keeping one
/// definition means trim and split agree on what whitespace is.
const WS_CLASS: &str = concat!(
    "\\t\\n\\x0B\\x0C\\r",
    " ",
    "\\u{00A0}\\u{1680}\\u{2000}-\\u{200A}",
    "\\u{2028}\\u{2029}\\u{202F}\\u{205F}\\u{3000}\\u{FEFF}",
);

/// Build a regex from `template`, substituting `\s` with the JS whitespace class
/// and `\S` with its negation.
fn js_ws_regex(template: &str) -> Regex {
    let positive = format!("[{WS_CLASS}]");
    let negative = format!("[^{WS_CLASS}]");
    let pattern = template.replace(r"\s", &positive).replace(r"\S", &negative);
    Regex::new(&pattern).expect("lexer regex compiles")
}

/// The active splitting regexes, in precedence order.
struct Regexes {
    urls: Regex,
    ids: Regex,
    time: Regex,
    number: Regex,
    space: Regex,
    email: Regex,
    punctuation: Regex,
    unblank: Regex,
}

fn regexes() -> &'static Regexes {
    static REGEXES: OnceLock<Regexes> = OnceLock::new();
    REGEXES.get_or_init(|| {
        // `(?i)` is case-insensitive. The URL pattern spells out ASCII character
        // classes instead of `\w` and `\d` so it matches ASCII word and digit
        // code points only, the same set JavaScript's ASCII classes cover.
        Regexes {
            urls: Regex::new(
                r"(?i)(?:https?://)(?:[0-9a-z.-]+)\.(?:[a-z.]{2,6})(?:[/0-9A-Za-z_.\-?#=&%]*)*/?",
            )
            .expect("urls regex compiles"),
            ids: js_ws_regex(r"(?i)(?:^|\s)([a-z0-9-]{8,45})"),
            time: js_ws_regex(r"(?i)(?:[0-9]|0[0-9]|1[0-9]|2[0-3]):(?:[0-5][0-9])\s?(?:[ap][m])"),
            number: Regex::new(r"(?i)[0-9]*\.[0-9]+|[0-9]+").expect("number regex compiles"),
            space: js_ws_regex(r"\s+"),
            email: Regex::new(
                r#"(?i)[-!#$%&'*+/0-9=?A-Z^_a-z{|}~](?:\.?[-!#$%&'*+/0-9=?A-Z^_a-z`{|}~])*@[a-zA-Z0-9](?:-?\.?[a-zA-Z0-9])*(?:\.[a-zA-Z](?:-?[a-zA-Z0-9])*)+"#,
            )
            .expect("email regex compiles"),
            punctuation: Regex::new(r#"[/.,?!"':;$()#’`]"#).expect("punctuation regex compiles"),
            unblank: js_ws_regex(r"\S"),
        }
    })
}

/// An English-aware tokenizer.
///
/// Construct one with [`Lexer::new`] and call [`Lexer::lex`]. The struct holds no
/// state, so a single instance is reusable and cheap to create.
///
/// ```
/// use pos_tagger::Lexer;
/// let lexer = Lexer::new();
/// assert_eq!(lexer.lex("I made $5.42 today."), vec!["I", "made", "$", "5.42", "today", "."]);
/// ```
#[derive(Debug, Default, Clone, Copy)]
pub struct Lexer;

impl Lexer {
    /// Create a lexer.
    pub fn new() -> Self {
        Lexer
    }

    /// Split `string` into tokens.
    ///
    /// Returns tokens in textual order. URLs, emails, long IDs, times, and
    /// numbers stay whole. Single punctuation marks split off as their own
    /// tokens. Whitespace-only fragments are dropped and every token is trimmed.
    pub fn lex(&self, string: &str) -> Vec<String> {
        let re = regexes();
        let order = [
            SplitRule::new(&re.urls),
            SplitRule::id(&re.ids),
            SplitRule::new(&re.time),
            SplitRule::new(&re.email),
            SplitRule::new(&re.number),
            SplitRule::new(&re.space),
            SplitRule::new(&re.punctuation),
        ];
        let root = LexerNode::new(string, order[0], &order[1..]);
        let mut out = Vec::new();
        root.fill_array(&mut out, &re.unblank);
        out
    }
}

#[derive(Clone, Copy)]
struct SplitRule<'a> {
    regex: &'a Regex,
    id_capture: bool,
}

impl<'a> SplitRule<'a> {
    fn new(regex: &'a Regex) -> Self {
        Self {
            regex,
            id_capture: false,
        }
    }

    fn id(regex: &'a Regex) -> Self {
        Self {
            regex,
            id_capture: true,
        }
    }
}

/// One node in the recursive split tree.
///
/// A node holds the delimiters its regex matched and its children. At the last
/// regex level children are raw gap strings. Above that they are child nodes.
enum LexerNode {
    /// Interior node: matched delimiters plus child nodes for each gap.
    Branch {
        matches: Vec<String>,
        children: Vec<LexerNode>,
    },
    /// Leaf node: matched delimiters plus raw gap strings.
    Leaf {
        matches: Vec<String>,
        gaps: Vec<String>,
    },
}

impl LexerNode {
    /// Build a node from `string`, splitting on `regex` and recursing on the
    /// remaining regexes.
    fn new(string: &str, rule: SplitRule<'_>, rest: &[SplitRule<'_>]) -> LexerNode {
        let string = js_trim(string);

        let (matches, gaps) = if string.is_empty() {
            (Vec::new(), vec![string])
        } else {
            let (matches, gaps) = split_with_rule(rule, &string);
            if matches.is_empty() {
                (Vec::new(), vec![string])
            } else {
                (matches, gaps)
            }
        };

        if rest.is_empty() {
            LexerNode::Leaf { matches, gaps }
        } else {
            let next = rest[0];
            let next_rest = &rest[1..];
            let children = gaps
                .iter()
                .map(|gap| LexerNode::new(gap, next, next_rest))
                .collect();
            LexerNode::Branch { matches, children }
        }
    }

    /// Flatten this node into `out`, interleaving gap then delimiter.
    ///
    /// For each child index `i`, emit the child's tokens, then if a delimiter
    /// `matches[i]` exists and is not blank, emit it trimmed. Blank children are
    /// dropped.
    fn fill_array(&self, out: &mut Vec<String>, unblank: &Regex) {
        match self {
            LexerNode::Branch { matches, children } => {
                for (i, child) in children.iter().enumerate() {
                    child.fill_array(out, unblank);
                    if i < matches.len() && unblank.is_match(&matches[i]) {
                        out.push(js_trim(&matches[i]));
                    }
                }
            }
            LexerNode::Leaf { matches, gaps } => {
                for (i, gap) in gaps.iter().enumerate() {
                    if unblank.is_match(gap) {
                        out.push(js_trim(gap));
                    }
                    if i < matches.len() && unblank.is_match(&matches[i]) {
                        out.push(js_trim(&matches[i]));
                    }
                }
            }
        }
    }
}

/// Split `s` with a rule, dropping delimiters from gaps.
///
/// Regular rules behave like JavaScript `String.split` with a non-capturing
/// pattern. The ID rule uses capture group 1 as the delimiter so trailing
/// punctuation remains in a gap.
fn split_with_rule(rule: SplitRule<'_>, s: &str) -> (Vec<String>, Vec<String>) {
    if rule.id_capture {
        split_id_matches(rule.regex, s)
    } else {
        let matches: Vec<String> = rule
            .regex
            .find_iter(s)
            .map(|m| m.as_str().to_string())
            .collect();
        let gaps = rule.regex.split(s).map(|piece| piece.to_string()).collect();
        (matches, gaps)
    }
}

fn split_id_matches(regex: &Regex, s: &str) -> (Vec<String>, Vec<String>) {
    let mut matches = Vec::new();
    let mut gaps = Vec::new();
    let mut last_end = 0;

    for capture in regex.captures_iter(s) {
        let Some(body) = capture.get(1) else {
            continue;
        };

        if !is_id_boundary(s, body.end()) {
            continue;
        }

        gaps.push(s[last_end..body.start()].to_string());
        matches.push(body.as_str().to_string());
        last_end = body.end();
    }

    if matches.is_empty() {
        (matches, Vec::new())
    } else {
        gaps.push(s[last_end..].to_string());
        (matches, gaps)
    }
}

fn is_id_boundary(s: &str, end: usize) -> bool {
    let mut chars = s[end..].chars();
    match chars.next() {
        None => true,
        Some(c) if is_js_whitespace(c) => true,
        Some('.') => !matches!(chars.next(), Some(next) if next.is_ascii_digit()),
        Some(c) => is_id_punctuation(c),
    }
}

fn is_id_punctuation(c: char) -> bool {
    matches!(
        c,
        '/' | ',' | '?' | '!' | '"' | '\'' | ':' | ';' | '$' | '(' | ')' | '#' | '’' | '`'
    )
}

/// Trim using the JavaScript `String.trim` whitespace set.
///
/// This covers ASCII whitespace plus the Unicode space separators, the line and
/// paragraph separators, the no-break space, and the byte order mark. For ASCII
/// input it matches `str::trim`.
fn js_trim(s: &str) -> String {
    s.trim_matches(is_js_whitespace).to_string()
}

/// True if `c` is whitespace under JavaScript `String.trim` rules.
fn is_js_whitespace(c: char) -> bool {
    matches!(
        c,
        '\u{0009}'
            | '\u{000A}'
            | '\u{000B}'
            | '\u{000C}'
            | '\u{000D}'
            | '\u{0020}'
            | '\u{00A0}'
            | '\u{1680}'
            | '\u{2000}'
            ..='\u{200A}'
                | '\u{2028}'
                | '\u{2029}'
                | '\u{202F}'
                | '\u{205F}'
                | '\u{3000}'
                | '\u{FEFF}'
    )
}