.rs

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


#[derive(Debug)]
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)]
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 }
    }
}


fn main() {
    #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
    enum AB { A, B }

    let contains_aba = NFA {
        initial: [0].into(),
        accepting: [3].into(),
        transition: [
            (0, [
                (AB::A, [0, 1].into()),
                (AB::B, [0].into()),
            ].into()),
            (1, [
                (AB::B, [2].into()),
            ].into()),
            (2, [
                (AB::A, [3].into()),
            ].into()),
            (3, [
                (AB::A, [3].into()),
                (AB::B, [3].into()),
            ].into()),
        ].into(),
    };

    let contains_aba_dfa = contains_aba.to_dfa().naturalize();
    let contains_aba_min = contains_aba_dfa.minimize();

    println!("{:?}", contains_aba);
    println!("{:?}", contains_aba_dfa);
    println!("{:?}", contains_aba_min);
}