arcweight 0.3.0

//! FST union algorithm.
//!
//! Constructs the union of two weighted finite-state transducers, creating
//! a new FST that accepts the combined language of both input FSTs.
//!
//! For FSTs $`T_1`$ and $`T_2`$:
//! $`L(T_1 \cup T_2) = L(T_1) \cup L(T_2)`$
//!
//! Weights for strings accepted by both FSTs combine via semiring addition:
//! $(T_1 \cup T_2)(x) = T_1(x) \oplus T_2(x)$.
//!
//! # Complexity
//!
//! - **Time:** $`O(|V_1| + |V_2| + |E_1| + |E_2|)`$ for copying both FSTs
//! - **Space:** $`O(|V_1| + |V_2| + |E_1| + |E_2| + 1)`$ for the combined structure
//!
//! # 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.

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst};
use crate::semiring::Semiring;
use crate::Result;

/// Computes the union of two FSTs, creating an FST that accepts both languages.
///
/// Creates a new FST that accepts any string accepted by either input FST:
/// $`L(T_1 \cup T_2) = L(T_1) \cup L(T_2)`$
///
/// Weights for strings in both languages combine via semiring addition:
/// $(T_1 \cup T_2)(x) = T_1(x) \oplus T_2(x)$.
///
/// # Arguments
///
/// * `fst1` - The first input FST
/// * `fst2` - The second input FST
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`Semiring`]
/// * `F1` - First FST type implementing [`Fst<W>`]
/// * `F2` - Second FST type implementing [`Fst<W>`]
/// * `M` - Output FST type implementing [`MutableFst<W>`] and [`Default`]
///
/// # Returns
///
/// A new FST accepting the union of the two input languages.
///
/// # Errors
///
/// Returns [`Error::Algorithm`](crate::Error::Algorithm) if either input FST is invalid.
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // FST 1: accepts "a"
/// let mut fst1 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst1.add_state();
/// let s1 = fst1.add_state();
/// fst1.set_start(s0);
/// fst1.set_final(s1, TropicalWeight::one());
/// fst1.add_arc(s0, Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1));
///
/// // FST 2: accepts "b"
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst2.add_state();
/// let s1 = fst2.add_state();
/// fst2.set_start(s0);
/// fst2.set_final(s1, TropicalWeight::one());
/// fst2.add_arc(s0, Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s1));
///
/// // Union accepts both "a" and "b"
/// let result: VectorFst<TropicalWeight> = union(&fst1, &fst2)?;
///
/// // New start state + states from both FSTs
/// assert_eq!(result.num_states(), 1 + fst1.num_states() + fst2.num_states());
/// # 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.
///
/// [`Semiring`]: crate::semiring::Semiring
pub fn union<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: Semiring,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
{
    let mut result = M::default();

    // create new start state
    let new_start = result.add_state();
    result.set_start(new_start);

    // copy first FST
    let mut state_map1 = vec![None; fst1.num_states()];
    for state in fst1.states() {
        let new_state = result.add_state();
        state_map1[state as usize] = Some(new_state);

        if let Some(weight) = fst1.final_weight(state) {
            result.set_final(new_state, weight.clone());
        }
    }

    // connect new start to fst1's start
    if let Some(start1) = fst1.start() {
        if let Some(mapped_start1) = state_map1[start1 as usize] {
            result.add_arc(new_start, Arc::epsilon(W::one(), mapped_start1));
        }
    }

    // add arcs from fst1
    for state in fst1.states() {
        if let Some(new_state) = state_map1[state as usize] {
            for arc in fst1.arcs(state) {
                if let Some(new_nextstate) = state_map1[arc.nextstate as usize] {
                    result.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    // copy second FST
    let mut state_map2 = vec![None; fst2.num_states()];
    for state in fst2.states() {
        let new_state = result.add_state();
        state_map2[state as usize] = Some(new_state);

        if let Some(weight) = fst2.final_weight(state) {
            result.set_final(new_state, weight.clone());
        }
    }

    // connect new start to fst2's start
    if let Some(start2) = fst2.start() {
        if let Some(mapped_start2) = state_map2[start2 as usize] {
            result.add_arc(new_start, Arc::epsilon(W::one(), mapped_start2));
        }
    }

    // add arcs from fst2
    for state in fst2.states() {
        if let Some(new_state) = state_map2[state as usize] {
            for arc in fst2.arcs(state) {
                if let Some(new_nextstate) = state_map2[arc.nextstate as usize] {
                    result.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_union_basic() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s1, TropicalWeight::one());
        fst1.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 2, TropicalWeight::new(2.0), t1));

        let unioned: VectorFst<TropicalWeight> = union(&fst1, &fst2).unwrap();

        // Should have states from both FSTs plus new start state
        assert!(unioned.num_states() >= fst1.num_states() + fst2.num_states());
        assert!(unioned.start().is_some());

        // Union might use epsilon transitions or restructure the FST
        // At minimum, should have preserved the original structure somehow
        assert!(unioned.num_states() >= 2); // Should have at least the states from both FSTs
    }

    #[test]
    fn test_union_empty() {
        let fst1 = VectorFst::<TropicalWeight>::new();
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let s0 = fst2.add_state();
        fst2.set_start(s0);
        fst2.set_final(s0, TropicalWeight::one());

        let unioned: VectorFst<TropicalWeight> = union(&fst1, &fst2).unwrap();

        // Should be equivalent to fst2
        assert!(unioned.start().is_some());
    }
}