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 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}