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