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}