Skip to main content

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    states: Vec<DfaState>,
9    state_index: BTreeMap<AtnConfigSet, usize>,
10}
11
12impl Dfa {
13    pub const fn new(atn_start_state: usize, decision: usize) -> Self {
14        Self {
15            decision,
16            atn_start_state,
17            states: Vec::new(),
18            state_index: BTreeMap::new(),
19        }
20    }
21
22    pub const fn decision(&self) -> usize {
23        self.decision
24    }
25
26    pub const fn atn_start_state(&self) -> usize {
27        self.atn_start_state
28    }
29
30    pub fn states(&self) -> &[DfaState] {
31        &self.states
32    }
33
34    /// Inserts a DFA state or returns the existing state number for an
35    /// equivalent ATN configuration set.
36    pub fn add_state(&mut self, mut state: DfaState) -> usize {
37        if let Some(existing) = self.state_index.get(&state.configs) {
38            return *existing;
39        }
40        let state_number = self.states.len();
41        state.state_number = state_number;
42        self.state_index.insert(state.configs.clone(), state_number);
43        self.states.push(state);
44        state_number
45    }
46
47    pub fn state_mut(&mut self, state_number: usize) -> Option<&mut DfaState> {
48        self.states.get_mut(state_number)
49    }
50}
51
52#[derive(Clone, Debug, Eq, PartialEq)]
53pub struct DfaState {
54    pub state_number: usize,
55    pub configs: AtnConfigSet,
56    pub edges: BTreeMap<i32, usize>,
57    pub is_accept_state: bool,
58    pub prediction: Option<usize>,
59    pub requires_full_context: bool,
60}
61
62impl DfaState {
63    pub const fn new(configs: AtnConfigSet) -> Self {
64        Self {
65            state_number: usize::MAX,
66            configs,
67            edges: BTreeMap::new(),
68            is_accept_state: false,
69            prediction: None,
70            requires_full_context: false,
71        }
72    }
73
74    pub fn add_edge(&mut self, symbol: i32, target_state: usize) {
75        self.edges.insert(symbol, target_state);
76    }
77
78    pub fn edge(&self, symbol: i32) -> Option<usize> {
79        self.edges.get(&symbol).copied()
80    }
81
82    pub const fn mark_accept(&mut self, prediction: usize) {
83        self.is_accept_state = true;
84        self.prediction = Some(prediction);
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::prediction::{AtnConfig, AtnConfigSet, PredictionContext};
92
93    #[test]
94    fn dfa_reuses_equal_config_sets() {
95        let mut configs = AtnConfigSet::new();
96        configs.add(AtnConfig::new(1, 1, PredictionContext::empty()));
97        let state = DfaState::new(configs.clone());
98        let mut dfa = Dfa::new(0, 0);
99        assert_eq!(dfa.add_state(state), 0);
100        assert_eq!(dfa.add_state(DfaState::new(configs)), 0);
101    }
102}