logicaffeine_compile/model/tokenizer/
regex.rs

1//! The regular expressions of the pre-tokenizers that llama.cpp runs through `std::regex` / `std::wregex`
2//! (`unicode_regex_split_stl`), matched with the semantics libstdc++ gives them: ECMAScript grammar, depth-first
3//! backtracking, the first match in priority order (left alternative first, greedy quantifiers longest first), and
4//! the `regex_iterator` walk over one span — an empty match is retried at its end as non-null and continuous, and
5//! every later search runs with `match_prev_avail`.
6//!
7//! A pattern runs over code units of one [`Width`]: [`Width::Narrow`] is llama.cpp's *collapsed* text (one byte per
8//! codepoint, a Unicode category folded to one byte, matched by `std::regex` over signed `char`), [`Width::Wide`]
9//! is its codepoint text (`std::wregex`). The character classes are libstdc++'s in the classic locale, measured
10//! against it for every code unit by `model_tokenizer_lock`: `\s` is `\t \n \v \f \r` and space, `\d` is ASCII
11//! digits, `\w` is ASCII alphanumerics and `_`, and `.` is every unit but `\n` and `\r` (and, wide, U+2028/U+2029).
12//!
13//! The matcher keeps its own backtrack stack, so a long input cannot exhaust the thread's stack. A quantified group
14//! that can match the empty string is refused at compile time: libstdc++ admits one empty iteration where
15//! ECMAScript admits none, and no pre-tokenizer pattern depends on either.
16
17use std::fmt;
18
19/// The code-unit width a pattern runs over.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum Width {
22    /// One byte per unit, compared as signed `char` the way `std::regex` compares them.
23    Narrow,
24    /// One codepoint per unit (`std::wregex`, 32-bit `wchar_t`).
25    Wide,
26}
27
28/// A pattern the compiler refuses, by name.
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct RegexError {
31    pub pattern: String,
32    pub at: usize,
33    pub why: &'static str,
34}
35
36impl fmt::Display for RegexError {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        write!(f, "the pre-tokenizer pattern {:?} is refused at unit {}: {}", self.pattern, self.at, self.why)
39    }
40}
41
42impl std::error::Error for RegexError {}
43
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45enum Named {
46    Space,
47    Digit,
48    Word,
49}
50
51impl Named {
52    fn contains(self, u: u32) -> bool {
53        match self {
54            Named::Space => matches!(u, 0x09..=0x0D | 0x20),
55            Named::Digit => (0x30..=0x39).contains(&u),
56            Named::Word => {
57                (0x30..=0x39).contains(&u) || (0x41..=0x5A).contains(&u) || (0x61..=0x7A).contains(&u) || u == 0x5F
58            }
59        }
60    }
61}
62
63/// A unit's position in range order: `std::regex` compares `char`s, which are signed.
64fn key(width: Width, u: u32) -> i64 {
65    match width {
66        Width::Narrow => u as u8 as i8 as i64,
67        Width::Wide => u as i64,
68    }
69}
70
71#[derive(Clone, Debug, PartialEq, Eq)]
72enum ClassItem {
73    Unit(u32),
74    Range(i64, i64),
75    Named { class: Named, negate: bool },
76}
77
78#[derive(Clone, Debug, PartialEq, Eq)]
79struct Class {
80    negate: bool,
81    items: Vec<ClassItem>,
82}
83
84#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85enum Look {
86    LineStart,
87    LineEnd,
88    WordBoundary,
89    NotWordBoundary,
90}
91
92#[derive(Clone, Debug, PartialEq, Eq)]
93enum Node {
94    Unit(u32),
95    Class(Class),
96    Any,
97    Concat(Vec<Node>),
98    Alt(Vec<Node>),
99    Repeat { node: Box<Node>, min: u32, max: Option<u32>, greedy: bool },
100    Look(Look),
101    Ahead { node: Box<Node>, negate: bool },
102}
103
104impl Node {
105    /// Whether the node can match without consuming a unit.
106    fn nullable(&self) -> bool {
107        match self {
108            Node::Unit(_) | Node::Class(_) | Node::Any => false,
109            Node::Concat(v) => v.iter().all(Node::nullable),
110            Node::Alt(v) => v.iter().any(Node::nullable),
111            Node::Repeat { node, min, .. } => *min == 0 || node.nullable(),
112            Node::Look(_) | Node::Ahead { .. } => true,
113        }
114    }
115}
116
117struct Parser<'a> {
118    p: &'a [u32],
119    i: usize,
120    pattern: &'a str,
121    width: Width,
122}
123
124impl<'a> Parser<'a> {
125    fn err(&self, why: &'static str) -> RegexError {
126        RegexError { pattern: self.pattern.to_string(), at: self.i, why }
127    }
128
129    fn peek(&self) -> Option<u32> {
130        self.p.get(self.i).copied()
131    }
132
133    fn peek_at(&self, k: usize) -> Option<u32> {
134        self.p.get(self.i + k).copied()
135    }
136
137    fn eat(&mut self, c: char) -> bool {
138        if self.peek() == Some(c as u32) {
139            self.i += 1;
140            true
141        } else {
142            false
143        }
144    }
145
146    fn disjunction(&mut self) -> Result<Node, RegexError> {
147        let mut alts = vec![self.alternative()?];
148        while self.eat('|') {
149            alts.push(self.alternative()?);
150        }
151        Ok(if alts.len() == 1 { alts.pop().expect("one alternative") } else { Node::Alt(alts) })
152    }
153
154    fn alternative(&mut self) -> Result<Node, RegexError> {
155        let mut terms = Vec::new();
156        while let Some(c) = self.peek() {
157            if c == '|' as u32 || c == ')' as u32 {
158                break;
159            }
160            terms.push(self.term()?);
161        }
162        Ok(if terms.len() == 1 { terms.pop().expect("one term") } else { Node::Concat(terms) })
163    }
164
165    fn term(&mut self) -> Result<Node, RegexError> {
166        let c = self.peek().expect("a term starts at a unit");
167        if c == '^' as u32 {
168            self.i += 1;
169            return Ok(Node::Look(Look::LineStart));
170        }
171        if c == '$' as u32 {
172            self.i += 1;
173            return Ok(Node::Look(Look::LineEnd));
174        }
175        if c == '\\' as u32 {
176            match self.peek_at(1) {
177                Some(b) if b == 'b' as u32 => {
178                    self.i += 2;
179                    return Ok(Node::Look(Look::WordBoundary));
180                }
181                Some(b) if b == 'B' as u32 => {
182                    self.i += 2;
183                    return Ok(Node::Look(Look::NotWordBoundary));
184                }
185                _ => {}
186            }
187        }
188        if c == '(' as u32 && self.peek_at(1) == Some('?' as u32) {
189            match self.peek_at(2) {
190                Some(k) if k == '=' as u32 || k == '!' as u32 => {
191                    self.i += 3;
192                    let node = self.disjunction()?;
193                    if !self.eat(')') {
194                        return Err(self.err("a lookahead is not closed"));
195                    }
196                    return Ok(Node::Ahead { node: Box::new(node), negate: k == '!' as u32 });
197                }
198                _ => {}
199            }
200        }
201        let atom = self.atom()?;
202        self.quantified(atom)
203    }
204
205    fn quantified(&mut self, atom: Node) -> Result<Node, RegexError> {
206        let (min, max) = match self.peek() {
207            Some(c) if c == '*' as u32 => {
208                self.i += 1;
209                (0, None)
210            }
211            Some(c) if c == '+' as u32 => {
212                self.i += 1;
213                (1, None)
214            }
215            Some(c) if c == '?' as u32 => {
216                self.i += 1;
217                (0, Some(1))
218            }
219            Some(c) if c == '{' as u32 => {
220                self.i += 1;
221                let min = self.decimal().ok_or_else(|| self.err("an interval has no lower bound"))?;
222                let max = if self.eat(',') {
223                    if self.peek() == Some('}' as u32) {
224                        None
225                    } else {
226                        Some(self.decimal().ok_or_else(|| self.err("an interval's upper bound is not a number"))?)
227                    }
228                } else {
229                    Some(min)
230                };
231                if !self.eat('}') {
232                    return Err(self.err("an interval is not closed"));
233                }
234                if max.is_some_and(|m| m < min) {
235                    return Err(self.err("an interval's upper bound is below its lower bound"));
236                }
237                (min, max)
238            }
239            _ => return Ok(atom),
240        };
241        let greedy = !self.eat('?');
242        if atom.nullable() && max != Some(1) {
243            return Err(self.err("a quantified group that can match the empty string"));
244        }
245        Ok(Node::Repeat { node: Box::new(atom), min, max, greedy })
246    }
247
248    fn decimal(&mut self) -> Option<u32> {
249        let start = self.i;
250        let mut v: u32 = 0;
251        while let Some(c) = self.peek() {
252            if !(0x30..=0x39).contains(&c) {
253                break;
254            }
255            v = v.checked_mul(10)?.checked_add(c - 0x30)?;
256            self.i += 1;
257        }
258        (self.i > start).then_some(v)
259    }
260
261    fn atom(&mut self) -> Result<Node, RegexError> {
262        let c = self.peek().expect("an atom starts at a unit");
263        self.i += 1;
264        match char::from_u32(c) {
265            Some('.') => Ok(Node::Any),
266            Some('(') => {
267                if self.peek() == Some('?' as u32) {
268                    if self.peek_at(1) == Some(':' as u32) {
269                        self.i += 2;
270                    } else {
271                        return Err(self.err("a group modifier other than (?: (?= (?!"));
272                    }
273                }
274                let node = self.disjunction()?;
275                if !self.eat(')') {
276                    return Err(self.err("a group is not closed"));
277                }
278                Ok(node)
279            }
280            Some(')') => Err(self.err("an unmatched )")),
281            Some('[') => self.class(),
282            Some('*') | Some('+') | Some('?') | Some('{') => Err(self.err("a quantifier with nothing to repeat")),
283            Some('\\') => self.atom_escape(),
284            _ => Ok(Node::Unit(c)),
285        }
286    }
287
288    fn hex(&mut self, n: usize) -> Result<u32, RegexError> {
289        let mut v = 0u32;
290        for _ in 0..n {
291            let d = self
292                .peek()
293                .and_then(char::from_u32)
294                .and_then(|c| c.to_digit(16))
295                .ok_or_else(|| self.err("a hex escape is short"))?;
296            v = v * 16 + d;
297            self.i += 1;
298        }
299        Ok(v)
300    }
301
302    /// The unit a character escape denotes, or the named class it is.
303    fn escape(&mut self, in_class: bool) -> Result<Result<u32, (Named, bool)>, RegexError> {
304        let c = self.peek().ok_or_else(|| self.err("a pattern ends in a backslash"))?;
305        self.i += 1;
306        Ok(match char::from_u32(c) {
307            Some('d') => Err((Named::Digit, false)),
308            Some('D') => Err((Named::Digit, true)),
309            Some('s') => Err((Named::Space, false)),
310            Some('S') => Err((Named::Space, true)),
311            Some('w') => Err((Named::Word, false)),
312            Some('W') => Err((Named::Word, true)),
313            Some('f') => Ok(0x0C),
314            Some('n') => Ok(0x0A),
315            Some('r') => Ok(0x0D),
316            Some('t') => Ok(0x09),
317            Some('v') => Ok(0x0B),
318            Some('b') if in_class => Ok(0x08),
319            Some('0') => Ok(0),
320            Some('x') => Ok(self.hex(2)?),
321            Some('u') => Ok(self.hex(4)?),
322            Some('c') => {
323                let l = self
324                    .peek()
325                    .filter(|&l| (0x41..=0x5A).contains(&l) || (0x61..=0x7A).contains(&l))
326                    .ok_or_else(|| self.err("a control escape is not followed by a letter"))?;
327                self.i += 1;
328                Ok(l % 32)
329            }
330            Some(d) if d.is_ascii_digit() => return Err(self.err("a back-reference")),
331            _ => Ok(c),
332        })
333    }
334
335    fn atom_escape(&mut self) -> Result<Node, RegexError> {
336        Ok(match self.escape(false)? {
337            Ok(u) => Node::Unit(u),
338            Err((class, negate)) => {
339                Node::Class(Class { negate: false, items: vec![ClassItem::Named { class, negate }] })
340            }
341        })
342    }
343
344    fn class(&mut self) -> Result<Node, RegexError> {
345        let negate = self.eat('^');
346        let mut items = Vec::new();
347        if self.peek() == Some(']' as u32) {
348            return Err(self.err("an empty bracket expression"));
349        }
350        loop {
351            let c = self.peek().ok_or_else(|| self.err("a bracket expression is not closed"))?;
352            if c == ']' as u32 {
353                self.i += 1;
354                break;
355            }
356            let lo = self.class_atom()?;
357            let dash = self.peek() == Some('-' as u32) && self.peek_at(1).is_some_and(|n| n != ']' as u32);
358            match lo {
359                Ok(lo) if dash => {
360                    self.i += 1;
361                    let hi = match self.class_atom()? {
362                        Ok(hi) => hi,
363                        Err(_) => return Err(self.err("a range ends in a class escape")),
364                    };
365                    let (lo, hi) = (key(self.width, lo), key(self.width, hi));
366                    if lo > hi {
367                        return Err(self.err("a range whose start is above its end"));
368                    }
369                    items.push(ClassItem::Range(lo, hi));
370                }
371                Ok(u) => items.push(ClassItem::Unit(u)),
372                Err((class, n)) => items.push(ClassItem::Named { class, negate: n }),
373            }
374        }
375        Ok(Node::Class(Class { negate, items }))
376    }
377
378    fn class_atom(&mut self) -> Result<Result<u32, (Named, bool)>, RegexError> {
379        let c = self.peek().expect("a class atom starts at a unit");
380        self.i += 1;
381        if c == '\\' as u32 {
382            self.escape(true)
383        } else {
384            Ok(Ok(c))
385        }
386    }
387}
388
389#[derive(Clone, Debug)]
390enum Inst {
391    Unit(u32),
392    Class(usize),
393    Any,
394    Split(usize, usize),
395    Jmp(usize),
396    Look(Look),
397    Ahead { body: usize, negate: bool },
398    Match,
399}
400
401/// A compiled pre-tokenizer pattern.
402#[derive(Clone, Debug)]
403pub struct Regex {
404    insts: Vec<Inst>,
405    classes: Vec<Class>,
406    width: Width,
407}
408
409struct Compiler {
410    insts: Vec<Inst>,
411    classes: Vec<Class>,
412    pending: Vec<(usize, Node)>,
413}
414
415impl Compiler {
416    fn emit(&mut self, i: Inst) -> usize {
417        self.insts.push(i);
418        self.insts.len() - 1
419    }
420
421    fn node(&mut self, n: &Node) {
422        match n {
423            Node::Unit(u) => {
424                self.emit(Inst::Unit(*u));
425            }
426            Node::Class(c) => {
427                self.classes.push(c.clone());
428                let id = self.classes.len() - 1;
429                self.emit(Inst::Class(id));
430            }
431            Node::Any => {
432                self.emit(Inst::Any);
433            }
434            Node::Concat(v) => v.iter().for_each(|n| self.node(n)),
435            Node::Alt(v) => {
436                let mut exits = Vec::new();
437                for (k, alt) in v.iter().enumerate() {
438                    if k + 1 < v.len() {
439                        let split = self.emit(Inst::Split(0, 0));
440                        self.node(alt);
441                        exits.push(self.emit(Inst::Jmp(0)));
442                        let next = self.insts.len();
443                        self.insts[split] = Inst::Split(split + 1, next);
444                    } else {
445                        self.node(alt);
446                    }
447                }
448                let end = self.insts.len();
449                for e in exits {
450                    self.insts[e] = Inst::Jmp(end);
451                }
452            }
453            Node::Repeat { node, min, max, greedy } => {
454                for _ in 0..*min {
455                    self.node(node);
456                }
457                match max {
458                    None => {
459                        let split = self.emit(Inst::Split(0, 0));
460                        self.node(node);
461                        self.emit(Inst::Jmp(split));
462                        let out = self.insts.len();
463                        self.insts[split] = if *greedy { Inst::Split(split + 1, out) } else { Inst::Split(out, split + 1) };
464                    }
465                    Some(max) => {
466                        let mut splits = Vec::new();
467                        for _ in *min..*max {
468                            splits.push(self.emit(Inst::Split(0, 0)));
469                            self.node(node);
470                        }
471                        let out = self.insts.len();
472                        for s in splits {
473                            self.insts[s] = if *greedy { Inst::Split(s + 1, out) } else { Inst::Split(out, s + 1) };
474                        }
475                    }
476                }
477            }
478            Node::Look(l) => {
479                self.emit(Inst::Look(*l));
480            }
481            Node::Ahead { node, negate } => {
482                let at = self.emit(Inst::Ahead { body: 0, negate: *negate });
483                self.pending.push((at, (**node).clone()));
484            }
485        }
486    }
487}
488
489/// How one search is anchored: where its attempt began and the `regex_constants` flags it runs with.
490#[derive(Clone, Copy)]
491struct Ctx {
492    begin: usize,
493    prev_avail: bool,
494    not_null: bool,
495}
496
497impl Regex {
498    /// Compile `pattern` for code units of `width`.
499    pub fn new(pattern: &str, width: Width) -> Result<Regex, RegexError> {
500        let units: Vec<u32> = pattern.chars().map(|c| c as u32).collect();
501        Regex::from_units(&units, pattern, width)
502    }
503
504    /// Compile a pattern given as code units — the collapsed narrow patterns carry raw bytes ≥ 0x80.
505    pub fn from_units(units: &[u32], pattern: &str, width: Width) -> Result<Regex, RegexError> {
506        let mut p = Parser { p: units, i: 0, pattern, width };
507        let node = p.disjunction()?;
508        if p.i != units.len() {
509            return Err(p.err("an unmatched )"));
510        }
511        let mut c = Compiler { insts: Vec::new(), classes: Vec::new(), pending: Vec::new() };
512        c.node(&node);
513        c.emit(Inst::Match);
514        while let Some((at, body)) = c.pending.pop() {
515            let start = c.insts.len();
516            c.node(&body);
517            c.emit(Inst::Match);
518            if let Inst::Ahead { negate, .. } = c.insts[at] {
519                c.insts[at] = Inst::Ahead { body: start, negate };
520            }
521        }
522        Ok(Regex { insts: c.insts, classes: c.classes, width })
523    }
524
525    fn class_has(&self, id: usize, u: u32) -> bool {
526        let c = &self.classes[id];
527        let k = key(self.width, u);
528        let hit = c.items.iter().any(|item| match *item {
529            ClassItem::Unit(v) => v == u,
530            ClassItem::Range(lo, hi) => lo <= k && k <= hi,
531            ClassItem::Named { class, negate } => class.contains(u) != negate,
532        });
533        hit != c.negate
534    }
535
536    fn any(&self, u: u32) -> bool {
537        match self.width {
538            Width::Narrow => u != 0x0A && u != 0x0D,
539            Width::Wide => !matches!(u, 0x0A | 0x0D | 0x2028 | 0x2029),
540        }
541    }
542
543    fn look(&self, l: Look, text: &[u32], pos: usize, ctx: Ctx) -> bool {
544        let word = |u: u32| Named::Word.contains(u);
545        match l {
546            Look::LineStart => pos == ctx.begin && !ctx.prev_avail,
547            Look::LineEnd => pos == text.len(),
548            Look::WordBoundary | Look::NotWordBoundary => {
549                let left = (pos != ctx.begin || ctx.prev_avail) && pos > 0 && word(text[pos - 1]);
550                let right = pos < text.len() && word(text[pos]);
551                (left != right) == (l == Look::WordBoundary)
552            }
553        }
554    }
555
556    /// The end of the first match starting exactly at `pos` from instruction `pc`, in priority order.
557    fn run(&self, pc: usize, pos: usize, text: &[u32], ctx: Ctx) -> Option<usize> {
558        let mut stack = vec![(pc, pos)];
559        while let Some((mut pc, mut pos)) = stack.pop() {
560            loop {
561                match self.insts[pc] {
562                    Inst::Unit(u) => {
563                        if pos < text.len() && text[pos] == u {
564                            pc += 1;
565                            pos += 1;
566                        } else {
567                            break;
568                        }
569                    }
570                    Inst::Class(id) => {
571                        if pos < text.len() && self.class_has(id, text[pos]) {
572                            pc += 1;
573                            pos += 1;
574                        } else {
575                            break;
576                        }
577                    }
578                    Inst::Any => {
579                        if pos < text.len() && self.any(text[pos]) {
580                            pc += 1;
581                            pos += 1;
582                        } else {
583                            break;
584                        }
585                    }
586                    Inst::Split(a, b) => {
587                        stack.push((b, pos));
588                        pc = a;
589                    }
590                    Inst::Jmp(a) => pc = a,
591                    Inst::Look(l) => {
592                        if self.look(l, text, pos, ctx) {
593                            pc += 1;
594                        } else {
595                            break;
596                        }
597                    }
598                    Inst::Ahead { body, negate } => {
599                        let sub = Ctx { begin: pos, ..ctx };
600                        if self.run(body, pos, text, sub).is_some() != negate {
601                            pc += 1;
602                        } else {
603                            break;
604                        }
605                    }
606                    Inst::Match => {
607                        if ctx.not_null && pos == ctx.begin {
608                            break;
609                        }
610                        return Some(pos);
611                    }
612                }
613            }
614        }
615        None
616    }
617
618    /// `regex_search` over `text[start..]`: the first start position with a match, and the match.
619    fn search(&self, text: &[u32], start: usize, prev_avail: bool, continuous: bool, not_null: bool) -> Option<(usize, usize)> {
620        let mut s = start;
621        loop {
622            let ctx = Ctx { begin: s, prev_avail: prev_avail || s > start, not_null };
623            if let Some(e) = self.run(0, s, text, ctx) {
624                return Some((s, e));
625            }
626            if continuous || s == text.len() {
627                return None;
628            }
629            s += 1;
630        }
631    }
632
633    /// Every match `std::regex_iterator` yields over `text`, as `(position, length)`.
634    pub fn matches(&self, text: &[u32]) -> Vec<(usize, usize)> {
635        let mut out = Vec::new();
636        let mut m = self.search(text, 0, false, false, false);
637        while let Some((ms, me)) = m {
638            out.push((ms, me - ms));
639            let mut start = me;
640            if ms == me {
641                if me == text.len() {
642                    break;
643                }
644                if let Some(next) = self.search(text, me, true, true, true) {
645                    m = Some(next);
646                    continue;
647                }
648                start += 1;
649            }
650            m = self.search(text, start, true, false, false);
651        }
652        out
653    }
654}
logicaffeine_compile/model/tokenizer/regex.rs

logicaffeine_compile/model/tokenizer/
regex.rs
