anyxml-automata 0.1.3

use std::{
    collections::{BTreeMap, HashMap, HashSet, VecDeque},
    hash::{BuildHasher, Hash, RandomState},
    marker::PhantomData,
    sync::atomic::AtomicUsize,
};

use crate::{Atom, ast::ASTNode};

#[derive(Debug, PartialEq, Eq, Default)]
struct FAFragment<A: Atom> {
    is_accepted: bool,
    // (start, end, to)
    rule_map: Vec<(A, A, usize)>,
}

impl<A: Atom> PartialOrd for FAFragment<A> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<A: Atom> Ord for FAFragment<A> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        if self.is_accepted != other.is_accepted {
            self.is_accepted.cmp(&other.is_accepted)
        } else {
            self.rule_map
                .iter()
                .map(|(start, end, to)| (start, end, to))
                .cmp(
                    other
                        .rule_map
                        .iter()
                        .map(|(start, end, to)| (start, end, to)),
                )
        }
    }
}

static AUTOMATON_ID: AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct State<A: Atom> {
    index: usize,
    // Identify which automaton generated this state.
    fa_id: usize,
    _phantom: PhantomData<fn() -> A>,
}

impl<A: Atom> State<A> {
    /// Check whether this state is generated from the given DFA.
    pub fn generated_by(&self, dfa: &DFA<A>) -> bool {
        self.fa_id == dfa.id
    }
}

impl<A: Atom> std::hash::Hash for State<A> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.index.hash(state);
        self.fa_id.hash(state);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TransitionError {
    /// A state generated from a different FA or an invalid state generated by an internal error.
    InvalidState,
    /// Reached a reject state and there is no state that can be returned.
    InputIsRejected,
}

impl std::fmt::Display for TransitionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl std::error::Error for TransitionError {}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FAAssembleError {
    InvalidQuantifier,
}

impl std::fmt::Display for FAAssembleError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl std::error::Error for FAAssembleError {}

#[derive(Debug)]
pub struct NFA<A: Atom> {
    id: usize,
    initial_state: usize,
    states: Vec<FAFragment<A>>,
}

impl<A: Atom> NFA<A> {
    fn empty() -> Self {
        let id = AUTOMATON_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        Self {
            id,
            initial_state: 0,
            states: vec![FAFragment {
                is_accepted: true,
                rule_map: vec![],
            }],
        }
    }

    pub fn assemble(ast: Option<&ASTNode<A>>) -> Result<NFA<A>, FAAssembleError> {
        let Some(ast) = ast else {
            return Ok(Self::empty());
        };

        // ranges: (point, is_end)
        fn collect_character_ranges<A: Atom>(ast: &ASTNode<A>, ranges: &mut Vec<(A, bool)>) {
            match ast {
                ASTNode::Charcters {
                    start,
                    end,
                    negation: _,
                } => {
                    ranges.push((*start, false));
                    ranges.push((*end, true));
                }
                ASTNode::Catenation(front, back) => {
                    collect_character_ranges(front, ranges);
                    collect_character_ranges(back, ranges);
                }
                ASTNode::Alternation(left, right) => {
                    collect_character_ranges(left, ranges);
                    collect_character_ranges(right, ranges);
                }
                ASTNode::ZeroOrOne(node) => {
                    collect_character_ranges(node, ranges);
                }
                ASTNode::ZeroOrMore(node) => {
                    collect_character_ranges(node, ranges);
                }
                ASTNode::OneOrMore(node) => {
                    collect_character_ranges(node, ranges);
                }
                ASTNode::Repeat {
                    node,
                    at_least: _,
                    at_most: _,
                } => {
                    collect_character_ranges(node, ranges);
                }
                ASTNode::RepeatExact(node, _) => {
                    collect_character_ranges(node, ranges);
                }
            }
        }

        let mut ranges = vec![(A::MIN, false), (A::MAX, true)];
        collect_character_ranges(ast, &mut ranges);
        ranges.sort_unstable();
        ranges.dedup();
        let ranges = ranges
            .windows(2)
            .filter_map(|v| {
                let (mut start, f) = v[0];
                let (mut end, g) = v[1];
                if start == end {
                    Some((start, end))
                } else {
                    if f {
                        start = start.next()?;
                    }
                    if !g {
                        end = end.previous()?;
                    }
                    (start <= end).then_some((start, end))
                }
            })
            .collect::<Vec<_>>();

        fn build_states<A: Atom>(
            ast: &ASTNode<A>,
            ranges: &[(A, A)],
            states: &mut Vec<FAFragment<A>>,
        ) -> Result<(usize, usize), FAAssembleError> {
            match ast {
                &ASTNode::Charcters {
                    start,
                    end,
                    negation,
                } => {
                    let mut pos = ranges.partition_point(|p| p.0 < start);
                    let mut from = FAFragment::default();
                    let to = FAFragment {
                        is_accepted: true,
                        ..Default::default()
                    };
                    if negation {
                        for &(f, t) in ranges.iter().take(pos) {
                            from.rule_map.push((f, t, states.len() + 1));
                        }
                        while pos < ranges.len() && ranges[pos].1 <= end {
                            pos += 1;
                        }
                        for &(f, t) in ranges.iter().skip(pos) {
                            from.rule_map.push((f, t, states.len() + 1));
                        }
                    } else {
                        while pos < ranges.len() && ranges[pos].1 <= end {
                            let (s, e) = ranges[pos];
                            from.rule_map.push((s, e, states.len() + 1));
                            pos += 1;
                        }
                    }
                    states.push(from);
                    states.push(to);
                    Ok((states.len() - 2, states.len() - 1))
                }
                ASTNode::Catenation(front, back) => {
                    // [sf -> ... -> ef] -> [sb -> ... -> eb]
                    let (sf, ef) = build_states(front, ranges, states)?;
                    let (sb, eb) = build_states(back, ranges, states)?;
                    states[ef].rule_map.push((A::EPSILON, A::EPSILON, sb));
                    states[ef].is_accepted = false;
                    Ok((sf, eb))
                }
                ASTNode::Alternation(left, right) => {
                    //       ┌─> [sl -> ... -> el] ─┐
                    // start ┤                      ├─> end
                    //       └─> [sr -> ... -> er] ─┘
                    let (sl, el) = build_states(left, ranges, states)?;
                    let (sr, er) = build_states(right, ranges, states)?;
                    let mut start = FAFragment::default();
                    start.rule_map.push((A::EPSILON, A::EPSILON, sl));
                    start.rule_map.push((A::EPSILON, A::EPSILON, sr));
                    states.push(start);
                    let mut end = FAFragment::default();
                    let len = states.len();
                    states[el].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[el].is_accepted = false;
                    states[er].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[er].is_accepted = false;
                    end.is_accepted = true;
                    states.push(end);
                    Ok((states.len() - 2, states.len() - 1))
                }
                ASTNode::ZeroOrOne(node) => {
                    // [start -> ... -> end]
                    //   │               ^
                    //   └───────────────┘
                    let (start, end) = build_states(node, ranges, states)?;
                    states[start].rule_map.push((A::EPSILON, A::EPSILON, end));
                    Ok((start, end))
                }
                ASTNode::ZeroOrMore(node) => {
                    // [start -> ... -> end] -> new_end
                    //   │ ^             │         ^
                    //   │ └─────────────┘         │
                    //   └─────────────────────────┘
                    let (start, end) = build_states(node, ranges, states)?;
                    let mut new_end = FAFragment::default();
                    states[end].is_accepted = false;
                    new_end.is_accepted = true;
                    let len = states.len();
                    states[start].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[end].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[end].rule_map.push((A::EPSILON, A::EPSILON, start));
                    states.push(new_end);
                    Ok((start, states.len() - 1))
                }
                ASTNode::OneOrMore(node) => {
                    // [start -> ... -> end] -> [start2 -> ... -> end2] -> new_end
                    //                            │ ^               │         ^
                    //                            │ └───────────────┘         │
                    //                            └───────────────────────────┘
                    let (start, end) = build_states(node, ranges, states)?;
                    let (start2, end2) = build_states(node, ranges, states)?;
                    states[end].is_accepted = false;
                    states[end2].is_accepted = false;
                    states[end].rule_map.push((A::EPSILON, A::EPSILON, start2));
                    let new_end = FAFragment {
                        is_accepted: true,
                        ..Default::default()
                    };
                    let len = states.len();
                    states[start2].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[end2].rule_map.push((A::EPSILON, A::EPSILON, len));
                    states[end2].rule_map.push((A::EPSILON, A::EPSILON, start2));
                    states.push(new_end);
                    Ok((start, states.len() - 1))
                }
                ASTNode::Repeat {
                    node,
                    at_least,
                    at_most,
                } => {
                    let new_start = FAFragment::default();
                    states.push(new_start);
                    let new_start = states.len() - 1;
                    let new_end = FAFragment {
                        is_accepted: true,
                        ..Default::default()
                    };
                    states.push(new_end);
                    let new_end = states.len() - 1;

                    let mut start = new_start;
                    for _ in 0..*at_least {
                        let (new_start, new_end) = build_states(node, ranges, states)?;
                        states[start]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_start));
                        states[new_end].is_accepted = false;
                        start = new_end;
                    }
                    if let &Some(at_most) = at_most {
                        if *at_least > at_most {
                            return Err(FAAssembleError::InvalidQuantifier);
                        }
                        for _ in *at_least..at_most {
                            let (new_start, end) = build_states(node, ranges, states)?;
                            states[start]
                                .rule_map
                                .push((A::EPSILON, A::EPSILON, new_start));
                            states[start]
                                .rule_map
                                .push((A::EPSILON, A::EPSILON, new_end));
                            states[end].is_accepted = false;
                            start = end;
                        }
                        states[start]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_end));
                    } else {
                        let (new_start, end) = build_states(node, ranges, states)?;
                        states[start]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_start));
                        states[start]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_end));
                        states[end].rule_map.push((A::EPSILON, A::EPSILON, new_end));
                        states[end]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_start));
                        states[end].is_accepted = false;
                    }

                    Ok((new_start, new_end))
                }
                ASTNode::RepeatExact(node, n) => {
                    let new_start = FAFragment::default();
                    states.push(new_start);
                    let new_start = states.len() - 1;
                    let new_end = FAFragment {
                        is_accepted: true,
                        ..Default::default()
                    };
                    states.push(new_end);
                    let new_end = states.len() - 1;

                    let mut start = new_start;
                    for _ in 0..*n {
                        let (new_start, end) = build_states(node, ranges, states)?;
                        states[start]
                            .rule_map
                            .push((A::EPSILON, A::EPSILON, new_start));
                        states[end].is_accepted = false;
                        start = end;
                    }

                    states[start]
                        .rule_map
                        .push((A::EPSILON, A::EPSILON, new_end));

                    Ok((new_start, new_end))
                }
            }
        }

        let mut states = vec![];
        let (initial_state, _) = build_states(ast, &ranges, &mut states)?;
        states.iter_mut().for_each(|state| {
            state.rule_map.sort_unstable();
            state.rule_map.dedup();
        });

        Ok(NFA {
            id: AUTOMATON_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
            initial_state,
            states,
        })
    }

    /// Check whether all transitions in this NFA are deterministic.  \
    /// This method can be used to check for ambiguities in the element content model.
    ///
    /// If `true`, you can expect that converting to a DFA will not significantly increase
    /// the number of states.
    pub fn is_deterministic(&self) -> bool {
        let mut cache = BTreeMap::new();
        let mut set = HashSet::new();
        for i in 0..self.states.len() {
            cache.clear();
            set.insert(i);
            let mut nt = vec![i];
            while let Some(now) = nt.pop() {
                let state = &self.states[now];
                let pos = state
                    .rule_map
                    .partition_point(|rule| (rule.0, rule.1) < (A::EPSILON, A::EPSILON));
                for rule in state.rule_map[pos..]
                    .iter()
                    .take_while(|rule| (rule.0, rule.1) == (A::EPSILON, A::EPSILON))
                {
                    if set.insert(rule.2) {
                        nt.push(rule.2);
                    }
                }
            }

            for state in set.drain() {
                for rule in self.states[state]
                    .rule_map
                    .iter()
                    .filter(|rule| (rule.0, rule.1) != (A::EPSILON, A::EPSILON))
                {
                    if let Some(to) = cache.insert((rule.0, rule.1), rule.2)
                        && to != rule.2
                    {
                        return false;
                    }
                }
            }
        }

        true
    }

    /// Returns the initial state.
    pub fn initial_state(&self) -> State<A> {
        State {
            index: self.initial_state,
            fa_id: self.id,
            _phantom: PhantomData,
        }
    }

    /// Returns the destination state when `input` is entered,
    /// using `initial_states` as the initial states.
    ///
    /// `initial_state` does not need to be a state generated by [`NFA::initial_state`];
    /// it can also be a state generated by [`DFA::transition`], etc.  \
    /// For example, you can use this when you want to enter a chunked string little by little.
    pub fn transition(
        &self,
        initial_states: impl Iterator<Item = State<A>>,
        input: impl Iterator<Item = A>,
    ) -> Result<Vec<State<A>>, TransitionError> {
        let mut current = HashSet::new();
        let mut nt = VecDeque::new();
        for state in initial_states {
            if state.fa_id != self.id || self.states.len() <= state.index {
                return Err(TransitionError::InvalidState);
            }
            if current.insert(state) {
                nt.push_back(state.index);
            }
        }

        let resolve_epsilon = |set: &mut HashSet<State<A>>, nt: &mut VecDeque<usize>| {
            while let Some(now) = nt.pop_front() {
                for rule in &self.states[now].rule_map {
                    if (rule.0, rule.1) == (A::EPSILON, A::EPSILON)
                        && set.insert(State {
                            index: rule.2,
                            fa_id: self.id,
                            _phantom: PhantomData,
                        })
                    {
                        nt.push_back(rule.2);
                    }
                }
            }
        };
        resolve_epsilon(&mut current, &mut nt);

        for c in input {
            let mut next = HashSet::new();
            for cur in current {
                for rule in &self.states[cur.index].rule_map {
                    if (rule.0..=rule.1).contains(&c) {
                        next.insert(State {
                            index: rule.2,
                            fa_id: self.id,
                            _phantom: PhantomData,
                        });
                    }
                }
            }

            if next.is_empty() {
                return Err(TransitionError::InputIsRejected);
            }

            current = {
                nt.extend(next.iter().map(|state| state.index));
                resolve_epsilon(&mut next, &mut nt);
                next
            };
        }

        Ok(current.into_iter().collect())
    }

    /// Check if the given state is an accepted state.
    ///
    /// If the given state is not generated from this DFA, always return false.
    pub fn is_accepted(&self, state: State<A>) -> bool {
        state.fa_id == self.id && self.states[state.index].is_accepted
    }

    /// Convert this NFA to DFA.
    pub fn build_dfa(&self) -> DFA<A> {
        let mut states = vec![];
        let mut cur = HashSet::from([self.initial_state]);
        let mut stack = vec![self.initial_state];
        while let Some(now) = stack.pop() {
            for &(_, _, to) in self.states[now]
                .rule_map
                .iter()
                .take_while(|v| v.0 == A::EPSILON)
            {
                if cur.insert(to) {
                    stack.push(to);
                }
            }
        }
        let hash_state = RandomState::new();

        // I use Zobrist Hashing to improve the efficiency of set matching and memory efficiency.
        fn generate_hash(set: &HashSet<usize>, state: &RandomState) -> u64 {
            set.iter()
                .copied()
                .map(|id| state.hash_one(id))
                .fold(0, |s, v| s ^ v)
        }

        fn build_states<A: Atom>(
            hash: u64,
            cur: &HashSet<usize>,
            hash_state: &RandomState,
            old_states: &[FAFragment<A>],
            states: &mut Vec<FAFragment<A>>,
            memo: &mut HashMap<u64, usize>,
        ) -> usize {
            let state_index = states.len();
            memo.insert(hash, state_index);
            states.push(FAFragment::default());
            if cur.iter().any(|index| old_states[*index].is_accepted) {
                states[state_index].is_accepted = true;
            }

            // List transitions other than ε.
            let mut rules = vec![];
            for &c in cur.iter() {
                rules.extend(
                    old_states[c]
                        .rule_map
                        .iter()
                        .filter(|&v| v.0 != A::EPSILON)
                        .cloned(),
                );
            }
            rules.sort_unstable_by_key(|v| (v.0, v.2));
            rules.dedup();

            if rules.is_empty() {
                return state_index;
            }
            let mut rule = (rules[0].0, rules[0].1);
            let mut buf = vec![];
            let mut set = HashSet::new();
            let mut push_state = |rule: (A, A), buf: &mut Vec<usize>| {
                set.clear();
                set.extend(buf.iter().copied());
                // `buf` contains a set of states that can be reached from the current state with a single transition based on a given rule.
                // Furthermore, we list the sets that can be reached using only ε from these states.
                while let Some(now) = buf.pop() {
                    for to in old_states[now]
                        .rule_map
                        .iter()
                        .take_while(|v| v.0 == A::EPSILON)
                    {
                        if set.insert(to.2) {
                            buf.push(to.2);
                        }
                    }
                }
                let hash = generate_hash(&set, hash_state);
                if let Some(to) = memo.get(&hash) {
                    // If the listed set has already been created, add the rule to the existing set.
                    states[state_index].rule_map.push((rule.0, rule.1, *to));
                } else {
                    // Otherwise, continue searching as it is a newly discovered set.
                    let index = build_states(hash, &set, hash_state, old_states, states, memo);
                    states[state_index].rule_map.push((rule.0, rule.1, index));
                }
            };
            for (start, end, to) in rules {
                if rule == (start, end) {
                    buf.push(to);
                    continue;
                }

                if !buf.is_empty() {
                    push_state(rule, &mut buf);
                    // Here, the buffer must be emptied, but since it is emptied in `push_state`,
                    // no additional operations are necessary.
                    debug_assert!(buf.is_empty());
                }

                rule = (start, end);
                buf.push(to);
            }

            if !buf.is_empty() {
                push_state(rule, &mut buf);
            }
            state_index
        }

        build_states(
            generate_hash(&cur, &hash_state),
            &cur,
            &hash_state,
            &self.states,
            &mut states,
            &mut HashMap::new(),
        );

        let mut dfa = DFA {
            id: self.id,
            states,
        };
        dfa.minimize();
        dfa
    }
}

#[derive(Debug)]
pub struct DFA<A: Atom> {
    id: usize,
    // The rule_map for each fragment must be sorted.
    // This constraint allows us to search for the transition destination using binary search.
    states: Vec<FAFragment<A>>,
}

impl<A: Atom> DFA<A> {
    fn empty() -> Self {
        let id = AUTOMATON_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        Self {
            id,
            states: vec![FAFragment {
                is_accepted: true,
                rule_map: vec![],
            }],
        }
    }

    /// Build a DFA from the AST. If the build fails, return `Err`.
    ///
    /// If `None` is entered, this function always returns `Ok`.  
    /// In this case, the DFA has only one state that is both the input state and the accepting state.  
    /// For example, it is likely that the regular expression converted to AST would be an empty string.
    pub fn assemble(ast: Option<&ASTNode<A>>) -> Result<DFA<A>, FAAssembleError> {
        let Some(ast) = ast else {
            return Ok(DFA::empty());
        };

        Ok(NFA::assemble(Some(ast))?.build_dfa())
    }

    /// Returns the initial state.
    pub fn initial_state(&self) -> State<A> {
        State {
            index: 0,
            fa_id: self.id,
            _phantom: PhantomData,
        }
    }

    /// Returns the destination state when `input` is entered, using `initial_state` as the initial state.
    ///
    /// `initial_state` does not need to be a state generated by [`DFA::initial_state`]; it can also be a state generated by [`DFA::transition`], etc.  
    /// For example, you can use this when you want to enter a chunked string little by little.
    pub fn transition(
        &self,
        initial_state: State<A>,
        input: impl Iterator<Item = A>,
    ) -> Result<State<A>, TransitionError> {
        if initial_state.fa_id != self.id {
            return Err(TransitionError::InvalidState);
        }

        let mut state = initial_state.index;
        if self.states.len() <= state {
            return Err(TransitionError::InvalidState);
        }

        for c in input {
            debug_assert!(
                self.states[state]
                    .rule_map
                    .windows(2)
                    .all(|v| v[0].1 < v[1].0)
            );
            state = self.states[state]
                .rule_map
                .binary_search_by(|(start, end, _)| {
                    if &c < start {
                        std::cmp::Ordering::Greater
                    } else if end < &c {
                        std::cmp::Ordering::Less
                    } else {
                        std::cmp::Ordering::Equal
                    }
                })
                .map(|index| self.states[state].rule_map[index].2)
                .map_err(|_| TransitionError::InputIsRejected)?;
        }

        Ok(State {
            index: state,
            fa_id: self.id,
            _phantom: PhantomData,
        })
    }

    /// Check whether the entered string is accepted.
    pub fn is_match(&self, input: impl Iterator<Item = A>) -> bool {
        self.transition(self.initial_state(), input)
            .is_ok_and(|state| self.is_accepted(state))
    }

    /// Check if the given state is an accepted state.
    ///
    /// If the given state is not generated from this DFA, always return false.
    pub fn is_accepted(&self, state: State<A>) -> bool {
        state.fa_id == self.id && self.states[state.index].is_accepted
    }

    fn minimize(&mut self) {
        let mut replace_to = (0..self.states.len()).collect::<Vec<_>>();
        let mut index = (0..self.states.len()).collect::<Vec<_>>();
        loop {
            index.sort_unstable_by_key(|&index| (&self.states[index], index));
            let mut update = false;
            for v in index.windows(2) {
                let l = v[0];
                let r = v[1];

                if self.states[l] == self.states[r] {
                    // Always merge to the smaller index.
                    // As a result, the index of the initial state is kept at 0.
                    replace_to[r] = replace_to[l];
                    update = true;
                }
            }

            if !update {
                break;
            }

            for &index in &index {
                if index == replace_to[index] {
                    for (_, _, to) in self.states[index].rule_map.iter_mut() {
                        *to = replace_to[*to];
                    }
                }
            }
            index.retain(|&index| index == replace_to[index]);
        }

        index.sort_unstable();
        for (to, &from) in index.iter().enumerate() {
            self.states.swap(to, from);
            replace_to[from] = to;
        }
        self.states.truncate(index.len());
        for state in self.states.iter_mut() {
            for (_, _, to) in state.rule_map.iter_mut() {
                *to = replace_to[*to];
            }
        }
    }
}