metrale_grammar/earley/parser.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// EarleyParser — the grammar-matching engine.
4// Port of `class EarleyParser` from xgrammar `cpp/earley_parser.{h,cc}`.
5//
6// This file holds the parser struct, its working state, and the public
7// API: construction, `advance`, `pop_last_states` (rollback),
8// `is_completed`, `push_state_and_expand`, `reset`. The three Earley
9// operations live in sibling files:
10// predict.rs — Predict + rule-reference expansion
11// complete.rs — Complete
12// scan.rs — Scan + the byte/char-class/FSM advance helpers
13
14use std::sync::Arc;
15
16use super::prune::ProductivityTable;
17use super::queue::ProcessQueue;
18use super::state::{NO_PREV_INPUT_POS, ParserState, UNEXPANDED_RULE_START_SEQUENCE_ID};
19use crate::grammar::GrammarData;
20use crate::support::Compact2DArray;
21
22/// One `(referenced_rule_id, parent_state)` entry of the
23/// completable-states table. The C++ stores
24/// `Compact2DArray<pair<int32_t, ParserState>>`; the Rust port now uses
25/// the ported [`Compact2DArray`] (CSR) layout, matching it exactly.
26pub type CompletableEntry = (i32, ParserState);
27
28/// The Earley parser. Drives grammar matching: it keeps the set of
29/// Earley items, advances them byte-by-byte, and supports
30/// push/rollback of parser state for the matcher's token rollback.
31#[derive(Debug)]
32pub struct EarleyParser {
33 /// The optimized, FSM-accelerated grammar being parsed.
34 pub(crate) grammar: Arc<GrammarData>,
35
36 /// `is_completed[i]` records, after consuming `i` inputs, whether
37 /// the root rule has completed (stop token acceptable).
38 pub(crate) is_completed: Vec<bool>,
39
40 /// `completable[i]` is the list of completable parent states
41 /// recorded at input position `i`. Earley completion consults it.
42 /// CSR layout (one flat buffer + offsets) — contiguous, alloc-once,
43 /// O(1) row views, and rollback is a cheap `pop_rows` truncation.
44 pub(crate) completable: Compact2DArray<CompletableEntry>,
45
46 /// `scanable_history[i]` holds the scanable states reached after
47 /// consuming input `i-1`. The push/rollback unit of the parser.
48 /// CSR layout, same rationale as `completable`.
49 pub(crate) scanable_history: Compact2DArray<ParserState>,
50
51 /// Scratch list of states to append to the scanable history in the
52 /// advance currently in progress.
53 pub(crate) to_be_added: Vec<ParserState>,
54
55 /// Reusable scratch holding a snapshot of the latest scanable row
56 /// while it is scanned. Avoids a per-`advance` heap allocation: the
57 /// CSR row cannot be borrowed across the `&mut self` `scan` calls,
58 /// so its (small, `Copy`) states are copied here once and reused.
59 pub(crate) scan_scratch: Vec<ParserState>,
60
61 /// Reusable scratch for the parents read out of a `completable` CSR
62 /// row during `Complete` / right-recursion expansion. The row can't
63 /// be borrowed across the `&mut self` queue/`completable` mutations,
64 /// so the needed entries are copied here once per call and reused.
65 pub(crate) parent_scratch: Vec<CompletableEntry>,
66
67 /// The predict/complete processing queue (with visited set).
68 pub(crate) queue: ProcessQueue,
69
70 /// Set true within an advance round when the root rule completes.
71 pub(crate) accept_stop_token: bool,
72
73 /// True after a one-off probe state has been pushed (see
74 /// `push_one_state_to_check`); cleared on the next rollback.
75 pub(crate) stop_token_is_accepted: bool,
76
77 /// Per-rule co-accessibility bitsets driving ZapFormat-style dead-
78 /// state pruning (Tier 3a). Computed once from the grammar's FSM
79 /// topology; consulted with O(1) bitset lookups after every advance.
80 /// See `prune.rs`.
81 pub(crate) productivity: ProductivityTable,
82}
83
84impl EarleyParser {
85 /// The default root initial state for `grammar`: the root rule,
86 /// unexpanded, at the no-previous-input root position.
87 pub(crate) fn root_initial_state(grammar: &GrammarData) -> ParserState {
88 ParserState::new(
89 grammar.root_rule_id(),
90 UNEXPANDED_RULE_START_SEQUENCE_ID,
91 0,
92 NO_PREV_INPUT_POS,
93 0,
94 )
95 }
96
97 /// Construct a parser for `grammar`, seeding it with `initial_state`
98 /// (or the root state when `initial_state` is the invalid
99 /// sentinel). When `need_expand` is false the initial state is only
100 /// placed in the scanable history — no predict/complete is run.
101 ///
102 /// Panics if the grammar is not optimized (matches the C++
103 /// `XGRAMMAR_LOG(FATAL)` contract — the FSM-accelerated parser
104 /// requires `per_rule_fsms`).
105 pub fn new(grammar: Arc<GrammarData>, initial_state: ParserState, need_expand: bool) -> Self {
106 assert!(
107 grammar.optimized,
108 "EarleyParser requires an optimized grammar (run GrammarOptimizer first)"
109 );
110 let init = if initial_state.is_invalid() {
111 Self::root_initial_state(&grammar)
112 } else {
113 initial_state
114 };
115
116 let productivity = ProductivityTable::build(&grammar);
117 let mut parser = Self {
118 grammar,
119 is_completed: Vec::new(),
120 completable: Compact2DArray::new(),
121 scanable_history: Compact2DArray::new(),
122 to_be_added: Vec::new(),
123 scan_scratch: Vec::new(),
124 parent_scratch: Vec::new(),
125 queue: ProcessQueue::new(),
126 accept_stop_token: false,
127 stop_token_is_accepted: false,
128 productivity,
129 };
130
131 if need_expand {
132 parser.push_state_and_expand(init);
133 } else {
134 parser.completable.push_row(&[]);
135 parser.is_completed.push(false);
136 parser.scanable_history.push_row(&[init]);
137 }
138 parser
139 }
140
141 /// Build a parser seeded with the grammar's root rule.
142 pub fn from_grammar(grammar: Arc<GrammarData>) -> Self {
143 Self::new(grammar, ParserState::invalid(), true)
144 }
145
146 /// True if the root rule has completed at the current input
147 /// position — i.e. the stop token may be emitted now.
148 pub fn is_completed(&self) -> bool {
149 *self.is_completed.last().unwrap_or(&false)
150 }
151
152 /// The scanable states reached at the current input position.
153 pub fn latest_scanable_states(&self) -> &[ParserState] {
154 if self.scanable_history.is_empty() {
155 &[]
156 } else {
157 self.scanable_history.back()
158 }
159 }
160
161 /// Number of input positions currently recorded (history depth).
162 pub fn num_steps(&self) -> usize {
163 self.scanable_history.len() as usize
164 }
165
166 /// Advance every scanable state by input byte `ch`.
167 ///
168 /// Returns true if `ch` is accepted by at least one state. When
169 /// `ch` is rejected the parser state is left unchanged, so the
170 /// caller may safely try another byte.
171 pub fn advance(&mut self, ch: u8) -> bool {
172 debug_assert!(self.queue.is_empty(), "queue must be empty before advance");
173 self.queue.clear_visited();
174 self.to_be_added.clear();
175 self.accept_stop_token = false;
176
177 // Scan phase: every scanable state of the latest step. The CSR
178 // row can't be held across the `&mut self` `scan` calls, so copy
179 // its (small, `Copy`) states into the reusable scratch buffer.
180 self.scan_scratch.clear();
181 if !self.scanable_history.is_empty() {
182 self.scan_scratch
183 .extend_from_slice(self.scanable_history.back());
184 }
185 for i in 0..self.scan_scratch.len() {
186 self.scan(self.scan_scratch[i], ch);
187 }
188
189 // Rejected: nothing was produced — leave state untouched.
190 if self.queue.is_empty() && self.to_be_added.is_empty() {
191 return false;
192 }
193
194 // Predict / Complete until the queue drains.
195 self.completable.push_row(&[]);
196 while let Some(state) = self.queue.pop() {
197 let (scanable, completable) = self.predict(state);
198 if completable {
199 self.complete(state);
200 }
201 if scanable {
202 self.to_be_added.push(state);
203 }
204 }
205
206 self.is_completed.push(self.accept_stop_token);
207 // ZapFormat dead-state pruning: drop scanable states whose FSM
208 // node can never reach its rule's end node before they enter the
209 // history. Language-preserving — see `prune.rs`.
210 self.productivity.prune(&mut self.to_be_added);
211 let added = std::mem::take(&mut self.to_be_added);
212 self.scanable_history.push_row_owned(added);
213 true
214 }
215}