redfa 0.0.3

use crate::derivatives::Differentiable;
use bit_set::BitSet;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::iter;

/// A state in a DFA.
#[derive(Debug, Clone)]
pub struct State<T, V> {
    /// Labelled transitions.
    pub by_char: BTreeMap<T, u32>,
    /// The default transition (for values not in `by_char`).
    /// Note that `by_char` is assumed not to cover the entire alphabet (`T`).
    pub default: u32,
    /// A value associated with the state.
    pub value: V,
}

/// A deterministic finite automaton (DFA), over the alphabet `T`.
/// Each state is annotated with a value of type `V`.
/// State 0 is the starting state.
#[derive(Debug, Clone)]
pub struct Dfa<T, V> {
    /// The list of states.
    pub states: Vec<State<T, V>>,
}

pub trait Normalize {
    fn normalize(self) -> Self;
}

impl<R: Normalize> Normalize for Vec<R> {
    fn normalize(self) -> Self {
        self.into_iter().map(Normalize::normalize).collect()
    }
}

impl<T, V> Dfa<T, V> {
    /// Construct a DFA from a list of differentiable objects.
    /// The elements of `initial` form the first states of the DFA.
    /// Returns the DFA, together with a mapping from derivatives to state numbers.
    pub fn from_derivatives(initial: Vec<V>) -> (Dfa<T, V>, BTreeMap<V, u32>)
    where
        T: Ord,
        V: Differentiable<T> + Normalize + Ord + Clone,
    {
        fn index<V: Ord + Clone>(
            &mut (ref mut indices, ref mut next): &mut (BTreeMap<V, u32>, VecDeque<V>),
            re: V,
        ) -> u32 {
            let next_index = indices.len() as u32;
            *indices
                .entry(re.clone()) // FIXME: unnecessary allocation
                .or_insert_with(|| {
                    next.push_back(re);
                    next_index
                })
        }

        let mut result = Dfa { states: Vec::new() };
        let mut worklist = (BTreeMap::new(), VecDeque::new());

        for r in initial {
            index(&mut worklist, r.normalize());
        }

        while let Some(re) = worklist.1.pop_front() {
            let d = re.derivative();
            let mut by_char = BTreeMap::new();
            let default = index(&mut worklist, d.rest.normalize());
            for (chars, dre) in d.d {
                let ix = index(&mut worklist, dre.normalize());
                // no point putting entries that are equal to the default
                if ix != default {
                    for ch in chars {
                        by_char.insert(ch, ix);
                    }
                }
            }
            result.states.push(State {
                by_char,
                default,
                value: re,
            });
        }

        (result, worklist.0)
    }

    /// Apply a function to each state's value.
    pub fn map<U, F>(self, mut f: F) -> Dfa<T, U>
    where
        F: FnMut(V) -> U,
    {
        Dfa {
            states: self
                .states
                .into_iter()
                .map(|state| State {
                    by_char: state.by_char,
                    default: state.default,
                    value: f(state.value),
                })
                .collect(),
        }
    }

    /// Find the reverse transitions from each state in the DFA.
    #[allow(clippy::type_complexity)]
    pub fn reverse(&self) -> Vec<(BTreeMap<&T, BTreeSet<u32>>, BTreeSet<u32>)>
    where
        T: Ord,
    {
        let mut result = vec![(BTreeMap::new(), BTreeSet::new()); self.states.len()];
        for (state_ix, state) in self.states.iter().enumerate() {
            let state_ix = state_ix as u32;
            let mut rev: BTreeMap<u32, BTreeSet<_>> = BTreeMap::new();
            for (by, &to) in &state.by_char {
                rev.entry(to).or_insert_with(BTreeSet::new).insert(by);
            }
            for (&to, by) in &rev {
                for what in by {
                    let (trans, default) = &mut result[to as usize];
                    trans
                        .entry(*what)
                        .or_insert_with(|| default.clone())
                        .insert(state_ix);
                }
            }
            // `state.default` means that, for all characters NOT in
            // `state.by_char`, there is a transition to `state_ix`.
            let &mut (ref mut trans, ref mut default) = &mut result[state.default as usize];
            for c in state.by_char.keys() {
                // make sure that characters in `state.by_char` are _excluded_
                trans.entry(c).or_insert_with(|| default.clone());
            }
            // any other characters, _not_ in `state.by_char`, should be
            // included
            for (key, ref mut val) in trans {
                if !state.by_char.contains_key(key) {
                    val.insert(state_ix);
                }
            }
            // as should the default
            default.insert(state_ix);
        }
        result
    }

    /// Minimize a DFA; i.e. find a DFA with the fewest states that is
    /// equivalent to the given DFA.
    /// Two DFAs are equivalent if, given the same string, they always lead to a
    /// state with the same associated value.
    pub fn minimize(&self) -> Dfa<T, &V>
    where
        T: Ord + Clone,
        V: Ord,
    {
        assert!(!self.states.is_empty());

        // `partitions` is a partition of the DFA states, representing the
        // current set of equivalence classes
        let mut partitions: Vec<BTreeSet<u32>> = vec![];
        {
            // Calculate the initial partition. The choice of initial partition
            // determines what states the algorithm considers distinguishable.
            // We only consider states that are reachable from the starting
            // state, so we traverse the DFA to find these states.
            let mut initial_partition = BTreeMap::new();
            let mut worklist = VecDeque::new();
            let mut seen = BitSet::new();
            worklist.push_back(0);
            seen.insert(0);
            while let Some(state_ix) = worklist.pop_front() {
                let state = &self.states[state_ix as usize];
                let part = *initial_partition.entry(&state.value).or_insert_with(|| {
                    let ix = partitions.len();
                    partitions.push(BTreeSet::new());
                    ix
                });
                partitions[part].insert(state_ix);
                for &next in state.by_char.values().chain(iter::once(&state.default)) {
                    if seen.insert(next as usize) {
                        worklist.push_back(next);
                    }
                }
            }
        }
        // The above code should always put state 0 in partition 0.
        debug_assert!(partitions[0].contains(&0));

        let preimages = self.reverse();
        let mut worklist: BTreeSet<usize> = (0..partitions.len()).collect();
        while let Some(&cur_ix) = worklist.iter().next() {
            // XXX: I wish there were a way to just grab the first element...
            worklist.remove(&cur_ix);
            let part = partitions[cur_ix].clone();
            let chars: BTreeSet<&T> = part
                .iter()
                .flat_map(|&state| preimages[state as usize].0.keys().copied())
                .collect();
            for c in chars.into_iter().map(Some).chain(iter::once(None)) {
                let mut l = BTreeSet::new();
                if let Some(c) = c {
                    for &state in &part {
                        if let Some(prevs) = preimages[state as usize].0.get(c) {
                            l.extend(prevs.iter().copied());
                        } else {
                            l.extend(preimages[state as usize].1.iter().copied());
                        }
                    }
                } else {
                    for &state in &part {
                        l.extend(preimages[state as usize].1.iter().copied());
                    }
                }
                let l = l;
                for part_ix in 0..partitions.len() {
                    let r1: BTreeSet<_> = partitions[part_ix].intersection(&l).copied().collect();
                    if r1.is_empty() {
                        continue;
                    }
                    let r2: BTreeSet<_> = partitions[part_ix].difference(&r1).copied().collect();
                    if r2.is_empty() {
                        continue;
                    }
                    // make sure that the starting state (#0) stays where it is
                    let (r1, r2) = if r2.contains(&0) { (r2, r1) } else { (r1, r2) };
                    // partitions[part_ix] = r1, partitions[new_ix] = r2
                    let new_ix = partitions.len();
                    // first update the worklist
                    if worklist.contains(&part_ix) {
                        // if the refined partition was already there, then keep
                        // both halves
                        worklist.insert(new_ix);
                    } else {
                        // otherwise, we need to add one half to the worklist; pick the smaller one
                        if r1.len() <= r2.len() {
                            worklist.insert(part_ix);
                        } else {
                            worklist.insert(new_ix);
                        }
                    }
                    // then refine partitions[part_ix]
                    partitions[part_ix] = r1;
                    partitions.push(r2);
                }
            }
        }

        // After refinement, the first partition should still contain the
        // starting state.
        debug_assert!(partitions[0].contains(&0));

        let partition: BTreeMap<u32, u32> = partitions
            .iter()
            .enumerate()
            .flat_map(|(part_ix, part)| {
                part.iter().map(move |&state_ix| (state_ix, part_ix as u32))
            })
            .collect();

        let states: Vec<_> = partitions
            .iter()
            .map(|part| {
                let state_ix = *part.iter().next().unwrap();
                let state = &self.states[state_ix as usize];
                let default = partition[&state.default];
                State {
                    by_char: state
                        .by_char
                        .iter()
                        .map(|(key, &to)| (key.clone(), partition[&to]))
                        .filter(|&(_, to)| to != default)
                        .collect(),
                    default,
                    value: &state.value,
                }
            })
            .collect();

        // Sanity check: the partitions should really be valid.
        if cfg!(debug_assertions) {
            for (part, new_state) in partitions.iter().zip(&states) {
                for &state_ix in part {
                    let state = &self.states[state_ix as usize];
                    for (key, &to) in &state.by_char {
                        assert_eq!(
                            partition[&to],
                            new_state
                                .by_char
                                .get(key)
                                .copied()
                                .unwrap_or(new_state.default)
                        );
                    }
                    assert!(partition[&state.default] == new_state.default);
                    assert!(state.value == *new_state.value);
                }
            }
        }

        Dfa { states }
    }
}

/// Compare DFAs by graph isomorphism.
impl<T: Ord, U, V: PartialEq<U>> PartialEq<Dfa<T, U>> for Dfa<T, V> {
    fn eq(&self, other: &Dfa<T, U>) -> bool {
        if self.states.len() != other.states.len() {
            return false;
        }
        let mut mapping = vec![None; self.states.len()];
        let mut worklist = VecDeque::new();
        mapping[0] = Some(0);
        worklist.push_back(0);
        while let Some(ix) = worklist.pop_front() {
            let other_ix = mapping[ix as usize].expect("worklist only contains populated entries");
            let a = &self.states[ix as usize];
            let b = &other.states[other_ix as usize];
            if a.by_char.len() != b.by_char.len() {
                return false;
            }
            for (c, &to) in &a.by_char {
                if let Some(&other_to) = b.by_char.get(c) {
                    if let Some(old_mapping) = mapping[to as usize].replace(other_to) {
                        // make sure the replaced element was the same
                        if old_mapping != other_to {
                            return false;
                        }
                    } else {
                        // new mapping
                        worklist.push_back(to);
                    }
                } else {
                    return false;
                }
            }
            if let Some(old_mapping) = mapping[a.default as usize].replace(b.default) {
                if old_mapping != b.default {
                    return false;
                }
            } else {
                worklist.push_back(a.default);
            }
        }
        true
    }
}
impl<T: Ord, V: Eq> Eq for Dfa<T, V> {}

impl<T, V> Dfa<T, V> {
    /// Compare DFAs by language equality.
    pub fn equiv<U>(&self, other: &Dfa<T, U>) -> bool
    where
        T: Ord + Clone,
        U: Ord,
        V: Ord + PartialEq<U>,
    {
        self.minimize() == other.minimize()
    }
}