.rs

use std::collections::{BTreeMap, BTreeSet};


#[derive(Debug, Clone)]
struct DFA<State, Action> {
    initial: State,
    accepting: BTreeSet<State>,
    transition: BTreeMap<State, BTreeMap<Action, State>>,
}

impl<State: Ord, Action: Ord> DFA<State, Action> {
    fn run(&self, sequence: impl Iterator<Item = Action>) -> Option<&State> {
        let mut current = &self.initial;
        for action in sequence {
            current = self.transition.get(current)?.get(&action)?;
        }
        Some(current)
    }

    fn accepts(&self, sequence: impl Iterator<Item = Action>) -> bool {
        self.run(sequence)
            .map(|state| self.accepting.contains(state))
            .unwrap_or(false)
    }
}

impl<State: Clone + Ord, Action: Clone + Ord> DFA<State, Action> {
    fn naturalize(&self) -> DFA<usize, Action> {
        let mut mapping = BTreeMap::new();
        let mut remap = move |state: &State| {
            if let Some(&mapped) = mapping.get(state) {
                mapped
            } else {
                let remapped = mapping.len();
                mapping.insert(state.clone(), remapped);
                remapped
            }
        };
        DFA {
            initial: remap(&self.initial),
            accepting: self.accepting.iter().map(|state| remap(state)).collect(),
            transition: self.transition.iter()
                .map(|(from, table)| {
                    let remapped_table = table.iter()
                        .map(|(action, to)| {
                            (action.clone(), remap(to))
                        }).collect();
                    (remap(from), remapped_table)
                }).collect(),
        }
    }

    fn to_nfa(&self) -> NFA<State, Action> {
        NFA {
            initial: [self.initial.clone()].into(),
            accepting: self.accepting.clone(),
            transition: self.transition.iter()
                .map(|(from, table)| {
                    (from.clone(), table.iter().map(|(action, to)| {
                        (action.clone(), [to.clone()].into())
                    }).collect())
                }).collect(),
        }
    }

    fn minimize(&self) -> DFA<usize, Action> {
        self.naturalize()
            .to_nfa()
            .reverse()
            .to_dfa()
            .naturalize()
            .to_nfa()
            .reverse()
            .to_dfa()
            .naturalize()
    }
}


#[derive(Debug, Clone)]
struct NFA<State, Action> {
    initial: BTreeSet<State>,
    accepting: BTreeSet<State>,
    transition: BTreeMap<State, BTreeMap<Action, BTreeSet<State>>>,
}

impl<State: Clone + Ord, Action: Ord> NFA<State, Action> {
    fn transit(&self, action: &Action, states: &BTreeSet<State>) -> BTreeSet<State> {
        let mut new_states = BTreeSet::new();
        for state in states {
            if let Some(table) = self.transition.get(state) {
                if let Some(to) = table.get(action) {
                    new_states.append(&mut to.clone());
                }
            }
        }
        new_states
    }

    fn run(&self, sequence: impl Iterator<Item = Action>) -> BTreeSet<State> {
        let mut current = self.initial.clone();
        for action in sequence {
            current = self.transit(&action, &current);
        }
        current
    }

    fn accepts(&self, sequence: impl Iterator<Item = Action>) -> bool {
        let states = self.run(sequence);
        states.iter().any(|state| self.accepting.contains(state))
    }
}

impl<State: Clone + Ord, Action: Clone + Ord> NFA<State, Action> {
    fn reverse(&self) -> NFA<State, Action> {
        let all_actions = self.transition
            .values()
            .flat_map(BTreeMap::keys)
            .cloned()
            .collect::<BTreeSet<Action>>();

        let transition = self.transition
            .keys()
            .map(|from| {
                let reverse_table = all_actions.iter()
                    .map(|action| {
                        let reverse_to = self.transition.iter()
                            .filter(|&(_, table)| {
                                table.get(action)
                                    .map(|from1| from1.contains(from))
                                    .unwrap_or(false)
                            })
                            .map(|(to, _)| to.clone())
                            .collect::<BTreeSet<State>>();
                        (action.clone(), reverse_to)
                    })
                    .filter(|(action, reverse_to)| !reverse_to.is_empty())
                    .collect();
                (from.clone(), reverse_table)
            })
            .collect();

        NFA {
            initial: self.accepting.clone(),
            accepting: self.initial.clone(),
            transition,
        }
    }

    fn to_dfa(&self) -> DFA<BTreeSet<State>, Action> {
        let all_actions = self.transition
            .values()
            .flat_map(BTreeMap::keys)
            .cloned()
            .collect::<BTreeSet<Action>>();

        let mut examined = BTreeSet::new();
        let mut unexamined = BTreeSet::from([self.initial.clone()]);
        while !unexamined.is_empty() {
            let mut new_unexamined = BTreeSet::new();
            for superstate in unexamined {
                for action in &all_actions {
                    let to = self.transit(action, &superstate);
                    if !examined.contains(&to) {
                        new_unexamined.insert(to);
                    }
                }
                examined.insert(superstate);
            }
            unexamined = new_unexamined;
        }

        let initial = self.initial.clone();

        let accepting = examined.iter()
            .filter(|superstate| superstate.iter().any(|state| self.accepting.contains(state)))
            .cloned()
            .collect();

        let transition = examined.iter()
            .map(|superstate| {
                let table = all_actions.iter()
                    .map(|action| (action.clone(), self.transit(action, superstate)))
                    .collect::<BTreeMap<Action, BTreeSet<State>>>();
                (superstate.clone(), table)
            })
            .collect();

        DFA { initial, accepting, transition }
    }
}


#[derive(Debug, Clone)]
enum RegExp<Action> {
    A(Action),
    Seq(Vec<RegExp<Action>>),
    Alt(Vec<RegExp<Action>>),
    Star(Box<RegExp<Action>>),
}

impl<Action: Clone + Ord> RegExp<Action> {
    fn to_nfa(&self, init: usize) -> (NFA<usize, Action>, usize) {
        let initial: BTreeSet<_> = [init+0].into();
        let mut accepting: BTreeSet<_> = [].into();
        let mut transition: BTreeMap<usize, BTreeMap<Action, BTreeSet<usize>>> = [].into();
        let mut size = init+1;

        let mut add_transitions = |from: usize, table: &BTreeMap<Action, BTreeSet<usize>>| {
            if let Some(old_table) = transition.get_mut(&from) {
                for (action, tos) in table {
                    if let Some(old_tos) = old_table.get_mut(action) {
                        old_tos.append(&mut tos.clone());
                    } else {
                        old_table.insert(action.clone(), tos.clone());
                    }
                }
            } else {
                transition.insert(from, table.clone());
            }
        };

        match self {
            RegExp::A(action) => {
                add_transitions(init+0, &[(action.clone(), [init+1].into())].into());
                accepting.insert(init+1);
                size += 1;
            }

            RegExp::Seq(regexps) => {
                accepting.insert(init+0);
                for regexp in regexps {
                    let (mut nfa, new_size) = regexp.to_nfa(size);
                    for (from, table) in nfa.transition {
                        if nfa.initial.contains(&from) {
                            for acc in &accepting {
                                add_transitions(*acc, &table);
                            }
                        }
                        add_transitions(from, &table);
                    }
                    if !nfa.accepting.iter().any(|acc| nfa.initial.contains(acc)) {
                        accepting.clear();
                    }
                    accepting.append(&mut nfa.accepting);
                    size = new_size;
                }
            }

            RegExp::Alt(regexps) => {
                for regexp in regexps {
                    let (mut nfa, new_size) = regexp.to_nfa(size);
                    for (from, table) in nfa.transition {
                        if nfa.initial.contains(&from) {
                            add_transitions(init+0, &table);
                        }
                        add_transitions(from, &table);
                    }
                    if nfa.accepting.iter().any(|acc| nfa.initial.contains(acc)) {
                        accepting.insert(init+0);
                    }
                    accepting.append(&mut nfa.accepting);
                    size = new_size;
                }
            }

            RegExp::Star(regexp) => {
                accepting.insert(init+0);
                let (mut nfa, new_size) = regexp.to_nfa(size);
                for (from, table) in nfa.transition {
                    add_transitions(from, &table);
                    if nfa.initial.contains(&from) {
                        add_transitions(init+0, &table);
                    }
                    for (action, tos) in table {
                        for to in tos {
                            if nfa.accepting.contains(&to) {
                                if nfa.initial.contains(&from) {
                                    add_transitions(init+0, &[(action.clone(), [init+0].into())].into());
                                }
                                add_transitions(from, &[(action.clone(), [init+0].into())].into());
                            }
                        }
                    }
                }
                size = new_size;
            }
        }

        (NFA { initial, accepting, transition }, size)
    }
}


fn main() {
    #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
    enum Event { Left, Right, Up, Down, Tick }

    let anything = RegExp::Star(Box::new(RegExp::Alt(vec![
        RegExp::A(Event::Left),
        RegExp::A(Event::Right),
        RegExp::A(Event::Up),
        RegExp::A(Event::Down),
        RegExp::A(Event::Tick),
    ])));

    let regexp = RegExp::Alt(vec![
        RegExp::Seq(vec![
            anything.clone(),
            RegExp::A(Event::Left),
            RegExp::Star(Box::new(RegExp::A(Event::Tick))),
        ]),
        RegExp::Seq(vec![
            anything.clone(),
            RegExp::A(Event::Right),
            RegExp::Star(Box::new(RegExp::A(Event::Tick))),
        ]),
        RegExp::Seq(vec![
            anything.clone(),
            RegExp::A(Event::Up),
            RegExp::Star(Box::new(RegExp::A(Event::Tick))),
        ]),
        RegExp::Seq(vec![
            anything.clone(),
            RegExp::A(Event::Down),
            RegExp::Star(Box::new(RegExp::A(Event::Tick))),
        ]),
        RegExp::Star(Box::new(RegExp::A(Event::Tick))),
    ]);

    let (nfa, size) = regexp.to_nfa(0);
    let dfa = nfa.to_dfa();
    let nat = dfa.naturalize();
    let min = dfa.minimize();

    println!("{:?}\n{}\n{:?}\n{:?}\n{:?}\n{:?}", regexp, size, nfa, dfa, nat, min);
}