arcweight 0.3.0

//! FST concatenation algorithm.
//!
//! Constructs the concatenation of two weighted finite-state transducers by
//! sequentially chaining them. The result FST accepts strings of the form $`xy`$
//! where $`x \in L(T_1)`$ and $`y \in L(T_2)`$.
//!
//! For FSTs $`T_1`$ and $`T_2`$:
//! $`L(T_1 \cdot T_2) = \{xy : x \in L(T_1), y \in L(T_2)\}`$
//!
//! # 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|)`$ 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;

/// Concatenates two FSTs to create an FST accepting $`T_1`$ followed by $`T_2`$.
///
/// Creates a new FST that accepts any string of the form $`xy`$ where $`x`$ is accepted
/// by the first FST and $`y`$ is accepted by the second FST:
/// $`L(T_1 \cdot T_2) = \{xy : x \in L(T_1), y \in L(T_2)\}`$
///
/// Weights combine through semiring multiplication: $(T_1 \cdot T_2)(xy) = T_1(x) \otimes T_2(y)$.
///
/// # Arguments
///
/// * `fst1` - The first FST (prefix)
/// * `fst2` - The second FST (suffix)
///
/// # 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 concatenation 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::new(0.5));
/// 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::new(0.3));
/// fst2.add_arc(s0, Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s1));
///
/// // Concatenation accepts "ab"
/// let result: VectorFst<TropicalWeight> = concat(&fst1, &fst2)?;
///
/// assert_eq!(result.num_states(), 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 concat<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();

    // 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());
        }
    }

    // set start state from fst1
    if let Some(start1) = fst1.start() {
        if let Some(new_start) = state_map1[start1 as usize] {
            result.set_start(new_start);
        }
    }

    // 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());
        }
    }

    // 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),
                    );
                }
            }
        }
    }

    // connect final states of fst1 to start of fst2
    if let Some(start2) = fst2.start() {
        if let Some(new_start2) = state_map2[start2 as usize] {
            for state in fst1.states() {
                if let Some(final_weight) = fst1.final_weight(state) {
                    if let Some(new_state) = state_map1[state as usize] {
                        // remove final weight
                        result.remove_final(new_state);
                        // add epsilon arc to start of fst2
                        result.add_arc(new_state, Arc::epsilon(final_weight.clone(), new_start2));
                    }
                }
            }
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_concat_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 concatenated: VectorFst<TropicalWeight> = concat(&fst1, &fst2).unwrap();

        // Should have states from both FSTs
        assert_eq!(
            concatenated.num_states(),
            fst1.num_states() + fst2.num_states()
        );
        assert!(concatenated.start().is_some());

        // Original final states of fst1 should no longer be final
        // Only final states of fst2 (with offset) should be final
        let mut final_count = 0;
        for state in concatenated.states() {
            if concatenated.is_final(state) {
                final_count += 1;
            }
        }
        assert!(final_count > 0);
    }

    #[test]
    fn test_concat_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 concatenated: VectorFst<TropicalWeight> = concat(&fst1, &fst2).unwrap();

        // Concat with empty should be empty
        assert!(concatenated.is_empty() || concatenated.num_arcs_total() == 0);
    }
}