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