arcweight 0.3.0

//! FST connection algorithm.
//!
//! Removes unreachable and non-productive states from weighted finite-state transducers,
//! ensuring all remaining states are both accessible from the start and can reach final states.
//!
//! The connection algorithm performs two reachability analyses: forward from the start state
//! to identify accessible states, and backward from final states to identify coaccessible
//! (productive) states. Only states that are both accessible and coaccessible are retained.
//!
//! # Complexity
//!
//! - **Time:** $`O(|V| + |E|)`$ where $`|V|`$ is the number of states and $`|E|`$ is the
//!   number of arcs
//! - **Space:** $`O(|V| + |E|)`$ for state tracking, reverse index, and result construction
//!
//! # References
//!
//! - John E. Hopcroft and Jeffrey D. Ullman. 1979. *Introduction to Automata Theory,
//!   Languages, and Computation*. Addison-Wesley, Reading, MA.
//! - Mehryar Mohri. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!   Automata*, Manfred Droste, Werner Kuich, and Heiko Vogler (Eds.). Springer,
//!   Berlin, Heidelberg, 213-254. DOI: <https://doi.org/10.1007/978-3-642-01492-5_6>

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::Semiring;
use crate::{Error, Result};
use rustc_hash::{FxHashMap, FxHashSet};

/// Removes non-accessible and non-coaccessible states from an FST.
///
/// Creates a new FST containing only states that are both accessible (reachable from
/// the start state) and coaccessible (can reach a final state). This operation removes
/// dead code and unreachable parts of the automaton while preserving the language.
///
/// For an FST T, a state q is:
/// - **Accessible:** There exists a path from start state to q
/// - **Coaccessible:** There exists a path from q to some final state
/// - **Useful:** Both accessible and coaccessible
///
/// The connected FST contains only useful states, eliminating unreachable
/// computation paths while maintaining the same accepted language:
/// $`L(\text{connect}(T)) = L(T)`$.
///
/// # Type Parameters
///
/// - `W` - The semiring weight type
/// - `F` - The input FST type implementing [`Fst<W>`]
/// - `M` - The output FST type implementing [`MutableFst<W>`] and [`Default`]
///
/// # Returns
///
/// A new FST containing only useful (accessible and coaccessible) states and their
/// connecting arcs.
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - The input FST has no start state
/// - Memory allocation fails during construction
///
/// # Panics
///
/// This function does not panic under normal operation.
///
/// # Examples
///
/// ## Removing Unreachable States
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // FST with unreachable state
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state(); // start state
/// let s1 = fst.add_state(); // reachable state
/// let s2 = fst.add_state(); // final state
/// let s3 = fst.add_state(); // unreachable isolated state
///
/// fst.set_start(s0);
/// fst.set_final(s2, TropicalWeight::one());
///
/// // Connected path: s0 -> s1 -> s2
/// fst.add_arc(s0, Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1));
/// fst.add_arc(s1, Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s2));
///
/// // s3 is isolated - no incoming or outgoing arcs
///
/// // Remove unreachable states
/// let connected: VectorFst<TropicalWeight> = connect(&fst)?;
///
/// // Result has 3 states (s0, s1, s2), s3 removed
/// assert_eq!(connected.num_states(), 3);
/// assert!(connected.num_states() < fst.num_states());
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// ## Removing Non-Coaccessible States
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // FST with dead-end state
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state(); // start
/// let s1 = fst.add_state(); // final state
/// let s2 = fst.add_state(); // dead-end state (not coaccessible)
///
/// fst.set_start(s0);
/// fst.set_final(s1, TropicalWeight::one());
///
/// // Good path: s0 -> s1 (final)
/// fst.add_arc(s0, Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1));
///
/// // Dead-end path: s0 -> s2 (but s2 cannot reach any final state)
/// fst.add_arc(s0, Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s2));
/// // s2 has no outgoing arcs to final states
///
/// // Remove non-coaccessible states
/// let connected: VectorFst<TropicalWeight> = connect(&fst)?;
///
/// // Result removes dead-end s2, keeps s0 and s1
/// assert_eq!(connected.num_states(), 2);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # References
///
/// - John E. Hopcroft and Jeffrey D. Ullman. 1979. *Introduction to Automata Theory,
///   Languages, and Computation*. Addison-Wesley, Reading, MA.
/// - Mehryar Mohri. 2009. Weighted automata algorithms. In *Handbook of Weighted
///   Automata*, Manfred Droste, Werner Kuich, and Heiko Vogler (Eds.). Springer,
///   Berlin, Heidelberg, 213-254. DOI: <https://doi.org/10.1007/978-3-642-01492-5_6>
pub fn connect<W, F, M>(fst: &F) -> Result<M>
where
    W: Semiring,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    let start = fst
        .start()
        .ok_or_else(|| Error::Algorithm("FST has no start state".into()))?;

    // find accessible states
    let accessible = find_accessible_states(fst, start);

    // find coaccessible states
    let coaccessible = find_coaccessible_states(fst);

    // keep only states that are both accessible and coaccessible
    let keep: FxHashSet<StateId> = accessible.intersection(&coaccessible).cloned().collect();

    if keep.is_empty() {
        return Ok(M::default());
    }

    // build new FST
    let mut result = M::default();
    let mut state_map = vec![None; fst.num_states()];

    // create new states
    for &state in &keep {
        let new_state = result.add_state();
        state_map[state as usize] = Some(new_state);
    }

    // set start
    if let Some(new_start) = state_map[start as usize] {
        result.set_start(new_start);
    }

    // copy arcs and final weights
    for &state in &keep {
        if let Some(new_state) = state_map[state as usize] {
            // final weight
            if let Some(weight) = fst.final_weight(state) {
                result.set_final(new_state, weight.clone());
            }

            // arcs
            for arc in fst.arcs(state) {
                if keep.contains(&arc.nextstate) {
                    if let Some(new_nextstate) = state_map[arc.nextstate as usize] {
                        result.add_arc(
                            new_state,
                            Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                        );
                    }
                }
            }
        }
    }

    Ok(result)
}

fn find_accessible_states<W: Semiring, F: Fst<W>>(fst: &F, start: StateId) -> FxHashSet<StateId> {
    let mut accessible = FxHashSet::default();
    let mut stack = vec![start];

    while let Some(state) = stack.pop() {
        if accessible.insert(state) {
            for arc in fst.arcs(state) {
                stack.push(arc.nextstate);
            }
        }
    }

    accessible
}

fn find_coaccessible_states<W: Semiring, F: Fst<W>>(fst: &F) -> FxHashSet<StateId> {
    // Build reverse arc index mapping each state to its predecessors: O(|V| + |E|)
    let mut predecessors: FxHashMap<StateId, Vec<StateId>> = FxHashMap::default();
    for state in fst.states() {
        for arc in fst.arcs(state) {
            predecessors.entry(arc.nextstate).or_default().push(state);
        }
    }

    let mut coaccessible = FxHashSet::default();
    let mut stack = Vec::new();

    // Start from all final states: O(|V|)
    for state in fst.states() {
        if fst.is_final(state) {
            stack.push(state);
        }
    }

    // Backward search using reverse index: O(|V| + |E|)
    while let Some(state) = stack.pop() {
        if coaccessible.insert(state) {
            // Process all predecessor states from reverse index
            if let Some(preds) = predecessors.get(&state) {
                for &pred in preds {
                    if !coaccessible.contains(&pred) {
                        stack.push(pred);
                    }
                }
            }
        }
    }

    coaccessible
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::prelude::*;
    use num_traits::One;

    #[test]
    fn test_connect_removes_unreachable() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();
        let s3 = fst.add_state(); // unreachable

        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::one());

        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s2, Arc::new(2, 2, TropicalWeight::new(1.0), s3)); // disconnected

        let connected: VectorFst<TropicalWeight> = connect(&fst).unwrap();

        // Should remove unreachable states
        assert!(connected.num_states() < fst.num_states());
        assert!(connected.start().is_some());
    }

    #[test]
    fn test_connect_already_connected() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();

        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::one());
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));

        let connected: VectorFst<TropicalWeight> = connect(&fst).unwrap();

        // Should be identical or similar
        assert_eq!(connected.num_states(), fst.num_states());
        assert!(connected.start().is_some());
    }
}