antlr4_runtime/
dfa.rs

1use crate::prediction::AtnConfigSet;
2use std::collections::BTreeMap;
3
4#[derive(Clone, Debug, Eq, PartialEq)]
5pub struct Dfa {
6    decision: usize,
7    atn_start_state: usize,
8    max_token_type: i32,
9    states: Vec<DfaState>,
10    state_index: BTreeMap<AtnConfigSet, usize>,
11    start_state: Option<usize>,
12    precedence_start_states: Vec<Option<usize>>,
13    precedence_mode: bool,
14    /// Set whenever a state, edge, or start state is learned. Lets the shared
15    /// decision-DFA cache skip cloning DFAs that a parse never extended, avoiding
16    /// per-parse churn when the cache is already warm. Not part of DFA identity.
17    dirty: bool,
18}
19
20impl Dfa {
21    pub const fn new(atn_start_state: usize, decision: usize) -> Self {
22        Self::with_max_token_type(atn_start_state, decision, 0)
23    }
24
25    pub const fn with_max_token_type(
26        atn_start_state: usize,
27        decision: usize,
28        max_token_type: i32,
29    ) -> Self {
30        Self {
31            decision,
32            atn_start_state,
33            max_token_type,
34            states: Vec::new(),
35            state_index: BTreeMap::new(),
36            start_state: None,
37            precedence_start_states: Vec::new(),
38            precedence_mode: false,
39            dirty: false,
40        }
41    }
42
43    pub const fn decision(&self) -> usize {
44        self.decision
45    }
46
47    pub const fn atn_start_state(&self) -> usize {
48        self.atn_start_state
49    }
50
51    pub const fn max_token_type(&self) -> i32 {
52        self.max_token_type
53    }
54
55    pub fn states(&self) -> &[DfaState] {
56        &self.states
57    }
58
59    pub const fn start_state(&self) -> Option<usize> {
60        self.start_state
61    }
62
63    pub const fn set_start_state(&mut self, state_number: usize) {
64        self.start_state = Some(state_number);
65        self.dirty = true;
66    }
67
68    /// Whether this DFA learned any new state/edge/start since it was created or
69    /// last cleared. The shared-cache merge uses this to skip untouched DFAs.
70    pub const fn is_dirty(&self) -> bool {
71        self.dirty
72    }
73
74    /// Clears the dirty flag, marking the current contents as the clean baseline
75    /// (called after publishing to / cloning from the shared cache).
76    pub const fn clear_dirty(&mut self) {
77        self.dirty = false;
78    }
79
80    pub const fn is_precedence_dfa(&self) -> bool {
81        self.precedence_mode
82    }
83
84    pub fn set_precedence_dfa(&mut self, precedence_dfa: bool) {
85        if self.precedence_mode == precedence_dfa {
86            return;
87        }
88        self.states.clear();
89        self.state_index.clear();
90        self.start_state = None;
91        self.precedence_start_states.clear();
92        self.precedence_mode = precedence_dfa;
93        self.dirty = true;
94        if precedence_dfa {
95            self.start_state = Some(self.add_state(DfaState::new(AtnConfigSet::new())));
96        }
97    }
98
99    pub fn precedence_start_state(&self, precedence: usize) -> Option<usize> {
100        self.precedence_start_states
101            .get(precedence)
102            .and_then(|state| *state)
103    }
104
105    pub(crate) fn precedence_start_states(&self) -> &[Option<usize>] {
106        &self.precedence_start_states
107    }
108
109    pub fn set_precedence_start_state(&mut self, precedence: usize, state_number: usize) {
110        if precedence >= self.precedence_start_states.len() {
111            self.precedence_start_states.resize(precedence + 1, None);
112        }
113        self.precedence_start_states[precedence] = Some(state_number);
114        self.dirty = true;
115    }
116
117    /// Inserts a DFA state or returns the existing state number for an
118    /// equivalent ATN configuration set.
119    pub fn add_state(&mut self, state: DfaState) -> usize {
120        if let Some(existing) = self.state_number_for_configs(&state.configs) {
121            return existing;
122        }
123        self.insert_state(state)
124    }
125
126    pub(crate) fn insert_state(&mut self, mut state: DfaState) -> usize {
127        let state_number = self.states.len();
128        state.state_number = state_number;
129        state.ensure_edge_capacity(self.max_token_type);
130        let state_key = state.configs.clone();
131        state.configs.set_readonly(true);
132        self.state_index.insert(state_key, state_number);
133        self.states.push(state);
134        self.dirty = true;
135        state_number
136    }
137
138    pub(crate) fn state_number_for_configs(&self, configs: &AtnConfigSet) -> Option<usize> {
139        self.state_index.get(configs).copied()
140    }
141
142    pub fn state(&self, state_number: usize) -> Option<&DfaState> {
143        self.states.get(state_number)
144    }
145
146    pub fn state_mut(&mut self, state_number: usize) -> Option<&mut DfaState> {
147        // Handing out a mutable state (used to add learned edges) conservatively
148        // marks the DFA dirty so the shared-cache merge re-publishes it.
149        self.dirty = true;
150        self.states.get_mut(state_number)
151    }
152}
153
154#[derive(Clone, Debug, Eq, PartialEq)]
155pub struct DfaState {
156    pub state_number: usize,
157    pub configs: AtnConfigSet,
158    pub edges: Vec<Option<usize>>,
159    pub is_accept_state: bool,
160    pub prediction: Option<usize>,
161    pub requires_full_context: bool,
162    pub conflicting_alts: Vec<usize>,
163    /// Whether any config for the predicted alt carries a semantic context.
164    /// Precomputed once at accept time (mirrors Go's `DFAState.predicates`) so
165    /// warm DFA hits don't rescan `configs` on every prediction lookup. Only
166    /// meaningful when `prediction` is `Some`; `false` for non-accept states.
167    ///
168    /// Crate-private: this is an internal derived cache (a pure function of
169    /// `configs` + `prediction`), kept off the public `DfaState` contract so it
170    /// is neither a struct-literal source break nor a surprising participant in
171    /// the public type's identity.
172    pub(crate) has_semantic_context_for_alt: bool,
173}
174
175impl DfaState {
176    pub const fn new(configs: AtnConfigSet) -> Self {
177        Self {
178            state_number: usize::MAX,
179            configs,
180            edges: Vec::new(),
181            is_accept_state: false,
182            prediction: None,
183            requires_full_context: false,
184            conflicting_alts: Vec::new(),
185            has_semantic_context_for_alt: false,
186        }
187    }
188
189    pub fn add_edge(&mut self, symbol: i32, target_state: usize) {
190        let Some(index) = edge_index(symbol) else {
191            return;
192        };
193        if index >= self.edges.len() {
194            self.edges.resize(index + 1, None);
195        }
196        self.edges[index] = Some(target_state);
197    }
198
199    pub fn edge(&self, symbol: i32) -> Option<usize> {
200        edge_index(symbol)
201            .and_then(|index| self.edges.get(index))
202            .and_then(|state| *state)
203    }
204
205    pub const fn mark_accept(&mut self, prediction: usize) {
206        self.is_accept_state = true;
207        self.prediction = Some(prediction);
208    }
209
210    fn ensure_edge_capacity(&mut self, max_token_type: i32) {
211        let Ok(max) = usize::try_from(max_token_type) else {
212            return;
213        };
214        if self.edges.len() < max + 2 {
215            self.edges.resize(max + 2, None);
216        }
217    }
218}
219
220fn edge_index(symbol: i32) -> Option<usize> {
221    if symbol < -1 {
222        None
223    } else {
224        usize::try_from(symbol + 1).ok()
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::prediction::{AtnConfig, AtnConfigSet, PredictionContext};
232
233    #[test]
234    fn dfa_reuses_equal_config_sets() {
235        let mut configs = AtnConfigSet::new();
236        configs.add(AtnConfig::new(1, 1, PredictionContext::empty()));
237        let state = DfaState::new(configs.clone());
238        let mut dfa = Dfa::with_max_token_type(0, 0, 16);
239        assert_eq!(dfa.add_state(state), 0);
240        assert_eq!(dfa.add_state(DfaState::new(configs)), 0);
241    }
242
243    #[test]
244    fn dfa_edges_are_dense_by_token_type() {
245        let mut state = DfaState::new(AtnConfigSet::new());
246        state.add_edge(-1, 3);
247        state.add_edge(5, 7);
248
249        assert_eq!(state.edge(-1), Some(3));
250        assert_eq!(state.edge(5), Some(7));
251        assert_eq!(state.edge(4), None);
252    }
253
254    #[test]
255    fn precedence_dfa_tracks_start_states_by_precedence() {
256        let mut dfa = Dfa::new(10, 2);
257        dfa.set_precedence_dfa(true);
258        dfa.set_precedence_start_state(4, 9);
259
260        assert!(dfa.is_precedence_dfa());
261        assert_eq!(dfa.precedence_start_state(4), Some(9));
262        assert_eq!(dfa.precedence_start_state(3), None);
263    }
264}
antlr4_runtime/dfa.rs

antlr4_runtime/
dfa.rs
