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 fn set_precedence_start_state(&mut self, precedence: usize, state_number: usize) {
106        if precedence >= self.precedence_start_states.len() {
107            self.precedence_start_states.resize(precedence + 1, None);
108        }
109        self.precedence_start_states[precedence] = Some(state_number);
110        self.dirty = true;
111    }
112
113    /// Inserts a DFA state or returns the existing state number for an
114    /// equivalent ATN configuration set.
115    pub fn add_state(&mut self, state: DfaState) -> usize {
116        if let Some(existing) = self.state_number_for_configs(&state.configs) {
117            return existing;
118        }
119        self.insert_state(state)
120    }
121
122    pub(crate) fn insert_state(&mut self, mut state: DfaState) -> usize {
123        let state_number = self.states.len();
124        state.state_number = state_number;
125        state.ensure_edge_capacity(self.max_token_type);
126        let state_key = state.configs.clone();
127        state.configs.set_readonly(true);
128        self.state_index.insert(state_key, state_number);
129        self.states.push(state);
130        self.dirty = true;
131        state_number
132    }
133
134    pub(crate) fn state_number_for_configs(&self, configs: &AtnConfigSet) -> Option<usize> {
135        self.state_index.get(configs).copied()
136    }
137
138    pub fn state(&self, state_number: usize) -> Option<&DfaState> {
139        self.states.get(state_number)
140    }
141
142    pub fn state_mut(&mut self, state_number: usize) -> Option<&mut DfaState> {
143        // Handing out a mutable state (used to add learned edges) conservatively
144        // marks the DFA dirty so the shared-cache merge re-publishes it.
145        self.dirty = true;
146        self.states.get_mut(state_number)
147    }
148}
149
150#[derive(Clone, Debug, Eq, PartialEq)]
151pub struct DfaState {
152    pub state_number: usize,
153    pub configs: AtnConfigSet,
154    pub edges: Vec<Option<usize>>,
155    pub is_accept_state: bool,
156    pub prediction: Option<usize>,
157    pub requires_full_context: bool,
158    pub conflicting_alts: Vec<usize>,
159    /// Whether any config for the predicted alt carries a semantic context.
160    /// Precomputed once at accept time (mirrors Go's `DFAState.predicates`) so
161    /// warm DFA hits don't rescan `configs` on every prediction lookup. Only
162    /// meaningful when `prediction` is `Some`; `false` for non-accept states.
163    ///
164    /// Crate-private: this is an internal derived cache (a pure function of
165    /// `configs` + `prediction`), kept off the public `DfaState` contract so it
166    /// is neither a struct-literal source break nor a surprising participant in
167    /// the public type's identity.
168    pub(crate) has_semantic_context_for_alt: bool,
169}
170
171impl DfaState {
172    pub const fn new(configs: AtnConfigSet) -> Self {
173        Self {
174            state_number: usize::MAX,
175            configs,
176            edges: Vec::new(),
177            is_accept_state: false,
178            prediction: None,
179            requires_full_context: false,
180            conflicting_alts: Vec::new(),
181            has_semantic_context_for_alt: false,
182        }
183    }
184
185    pub fn add_edge(&mut self, symbol: i32, target_state: usize) {
186        let Some(index) = edge_index(symbol) else {
187            return;
188        };
189        if index >= self.edges.len() {
190            self.edges.resize(index + 1, None);
191        }
192        self.edges[index] = Some(target_state);
193    }
194
195    pub fn edge(&self, symbol: i32) -> Option<usize> {
196        edge_index(symbol)
197            .and_then(|index| self.edges.get(index))
198            .and_then(|state| *state)
199    }
200
201    pub const fn mark_accept(&mut self, prediction: usize) {
202        self.is_accept_state = true;
203        self.prediction = Some(prediction);
204    }
205
206    fn ensure_edge_capacity(&mut self, max_token_type: i32) {
207        let Ok(max) = usize::try_from(max_token_type) else {
208            return;
209        };
210        if self.edges.len() < max + 2 {
211            self.edges.resize(max + 2, None);
212        }
213    }
214}
215
216fn edge_index(symbol: i32) -> Option<usize> {
217    if symbol < -1 {
218        None
219    } else {
220        usize::try_from(symbol + 1).ok()
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227    use crate::prediction::{AtnConfig, AtnConfigSet, PredictionContext};
228
229    #[test]
230    fn dfa_reuses_equal_config_sets() {
231        let mut configs = AtnConfigSet::new();
232        configs.add(AtnConfig::new(1, 1, PredictionContext::empty()));
233        let state = DfaState::new(configs.clone());
234        let mut dfa = Dfa::with_max_token_type(0, 0, 16);
235        assert_eq!(dfa.add_state(state), 0);
236        assert_eq!(dfa.add_state(DfaState::new(configs)), 0);
237    }
238
239    #[test]
240    fn dfa_edges_are_dense_by_token_type() {
241        let mut state = DfaState::new(AtnConfigSet::new());
242        state.add_edge(-1, 3);
243        state.add_edge(5, 7);
244
245        assert_eq!(state.edge(-1), Some(3));
246        assert_eq!(state.edge(5), Some(7));
247        assert_eq!(state.edge(4), None);
248    }
249
250    #[test]
251    fn precedence_dfa_tracks_start_states_by_precedence() {
252        let mut dfa = Dfa::new(10, 2);
253        dfa.set_precedence_dfa(true);
254        dfa.set_precedence_start_state(4, 9);
255
256        assert!(dfa.is_precedence_dfa());
257        assert_eq!(dfa.precedence_start_state(4), Some(9));
258        assert_eq!(dfa.precedence_start_state(3), None);
259    }
260}
antlr4_runtime/dfa.rs

antlr4_runtime/
dfa.rs
