arcweight 0.3.0

//! Linear FST decoding algorithms.
//!
//! Extracts the unique accepting path from linear finite-state transducers,
//! which have exactly one path from the start state to a final state.
//!
//! # Overview
//!
//! A linear FST is a special case where the automaton encodes a single
//! input-output pair (string transduction). This module provides efficient
//! extraction of that path, including the state sequence, arc sequence,
//! and total weight.
//!
//! Linear FSTs arise naturally in several contexts:
//! - Results of shortest-path algorithms on deterministic FSTs
//! - String-to-string transducers encoding single mappings
//! - Lattice best-path extraction in speech recognition
//!
//! # Algorithm
//!
//! The decoding algorithm performs a deterministic traversal:
//! 1. Start from the initial state
//! 2. At each state, verify exactly one outgoing arc (or epsilon chain)
//! 3. Follow the unique path, accumulating weights
//! 4. Terminate at a final state with the complete path
//!
//! # Complexity
//!
//! - **Time:** $`O(|V|)`$ where |V| is the number of states in the path
//! - **Space:** $`O(|V|)`$ for storing the extracted path
//!
//! # Use Cases
//!
//! - **Shortest Path Extraction:** Decode result of `shortest_path` algorithm
//! - **String Transduction:** Extract input/output pair from linear transducer
//! - **Lattice Decoding:** Extract best hypothesis from speech recognition lattice
//! - **Alignment Recovery:** Recover alignment from sequence-to-sequence models
//!
//! # Examples
//!
//! ```
//! use arcweight::prelude::*;
//!
//! // Create a linear FST encoding "ab" -> "xy"
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! let s0 = fst.add_state();
//! let s1 = fst.add_state();
//! let s2 = fst.add_state();
//! fst.set_start(s0);
//! fst.set_final(s2, TropicalWeight::one());
//! fst.add_arc(s0, Arc::new('a' as u32, 'x' as u32, TropicalWeight::new(1.0), s1));
//! fst.add_arc(s1, Arc::new('b' as u32, 'y' as u32, TropicalWeight::new(2.0), s2));
//!
//! let path = decode_linear_fst(&fst)?;
//!
//! assert_eq!(path.states.len(), 3);  // s0, s1, s2
//! assert_eq!(path.arcs.len(), 2);    // two transitions
//! assert_eq!(path.input_labels(), vec!['a' as u32, 'b' as u32]);
//! assert_eq!(path.output_labels(), vec!['x' as u32, 'y' as u32]);
//! # Ok::<(), arcweight::Error>(())
//! ```
//!
//! # References
//!
//! - Mohri, M. (2009). Weighted automata algorithms. In *Handbook of Weighted
//!   Automata* (pp. 213-254). Springer. <https://doi.org/10.1007/978-3-642-01492-5_6>
//! - Mohri, M., Pereira, F., and Riley, M. (2008). Speech recognition with weighted
//!   finite-state transducers. In *Springer Handbook of Speech Processing*
//!   (pp. 559-584). Springer. <https://doi.org/10.1007/978-3-540-49127-9_28>

use crate::fst::Fst;
use crate::semiring::Semiring;
use crate::utils::FstPath;
use crate::Result;

/// Decodes a linear FST, extracting the single accepting path.
///
/// A linear FST has exactly one path from the start state to a final state.
/// This function extracts that path, returning the sequence of states, arcs,
/// and the total weight.
///
/// # Errors
///
/// Returns an error if:
/// - The FST is not linear (has multiple paths or cycles)
/// - The FST has no start state
/// - The FST has no accepting path
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Create a linear FST: "hello" -> "hi"
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// let s2 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s2, TropicalWeight::one());
/// fst.add_arc(s0, Arc::new('h' as u32, 'h' as u32, TropicalWeight::one(), s1));
/// fst.add_arc(s1, Arc::new('i' as u32, 'i' as u32, TropicalWeight::one(), s2));
///
/// let path = decode_linear_fst(&fst)?;
/// assert_eq!(path.states.len(), 3);
/// assert_eq!(path.arcs.len(), 2);
/// # Ok::<(), arcweight::Error>(())
/// ```
pub fn decode_linear_fst<W, F>(fst: &F) -> Result<FstPath<W>>
where
    W: Semiring + Clone,
    F: Fst<W>,
{
    let start = fst
        .start()
        .ok_or_else(|| crate::Error::Algorithm("FST has no start state".to_string()))?;

    let mut states = vec![start];
    let mut arcs = Vec::new();
    let mut current = start;
    let mut weight = W::one();
    let mut visited = std::collections::HashSet::new();
    let mut just_followed_epsilon = false;

    // Traverse the linear path
    loop {
        // Check for cycles (skip if we just followed an epsilon arc, as we already checked)
        if !just_followed_epsilon {
            if visited.contains(&current) {
                return Err(crate::Error::Algorithm(
                    "FST is not linear: contains cycles".to_string(),
                ));
            }
            visited.insert(current);
        }
        just_followed_epsilon = false;

        // Check if we've reached a final state
        if let Some(final_weight) = fst.final_weight(current) {
            weight = weight * final_weight.clone();
            return Ok(FstPath {
                states,
                arcs,
                weight,
                final_state: current,
            });
        }

        // Get outgoing arcs - handle epsilon transitions
        let arc_iter = fst.arcs(current);
        let mut non_epsilon_arcs: Vec<_> = arc_iter.collect();

        // Separate epsilon and non-epsilon arcs
        let epsilon_arcs: Vec<_> = non_epsilon_arcs
            .iter()
            .filter(|arc| arc.is_epsilon())
            .cloned()
            .collect();
        non_epsilon_arcs.retain(|arc| !arc.is_epsilon());

        // Check for multiple non-epsilon paths (non-linear)
        if non_epsilon_arcs.len() > 1 {
            return Err(crate::Error::Algorithm(
                "FST is not linear: multiple non-epsilon paths from state".to_string(),
            ));
        }

        // Handle epsilon transitions first (if any)
        // For a linear FST, there should be at most one epsilon arc from each state
        // (or epsilons must form a deterministic chain)
        if epsilon_arcs.len() > 1 {
            return Err(crate::Error::Algorithm(
                "FST is not linear: multiple epsilon paths from state".to_string(),
            ));
        }

        if let Some(epsilon_arc) = epsilon_arcs.first() {
            // Follow epsilon transition
            weight = weight.clone() * epsilon_arc.weight.clone();
            states.push(epsilon_arc.nextstate);
            arcs.push(epsilon_arc.clone());
            current = epsilon_arc.nextstate;

            // Check for cycles after following epsilon
            if visited.contains(&current) {
                return Err(crate::Error::Algorithm(
                    "FST is not linear: contains cycles".to_string(),
                ));
            }
            visited.insert(current);
            just_followed_epsilon = true;

            // Check if we reached a final state via epsilon
            if let Some(final_weight) = fst.final_weight(current) {
                weight = weight * final_weight.clone();
                return Ok(FstPath {
                    states,
                    arcs,
                    weight,
                    final_state: current,
                });
            }

            // Continue loop to process arcs from the new state
            continue;
        }

        // Check for dead end (no path to final state)
        let arc = non_epsilon_arcs.first().ok_or_else(|| {
            crate::Error::Algorithm("FST is not linear: no path to final state".to_string())
        })?;

        // Follow the non-epsilon arc
        weight = weight * arc.weight.clone();
        states.push(arc.nextstate);
        arcs.push(arc.clone());
        current = arc.nextstate;
    }
}

/// Decode linear FST and extract input labels
///
/// Convenience function that extracts just the input label sequence.
pub fn decode_linear_fst_input<W, F>(fst: &F) -> Result<Vec<u32>>
where
    W: Semiring + Clone,
    F: Fst<W>,
{
    let path = decode_linear_fst(fst)?;
    Ok(path.input_labels())
}

/// Decode linear FST and extract output labels
///
/// Convenience function that extracts just the output label sequence.
pub fn decode_linear_fst_output<W, F>(fst: &F) -> Result<Vec<u32>>
where
    W: Semiring + Clone,
    F: Fst<W>,
{
    let path = decode_linear_fst(fst)?;
    Ok(path.output_labels())
}

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

    #[test]
    fn test_decode_linear_fst() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::new(0.5));
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(2.0), s2));

        let path = decode_linear_fst(&fst).unwrap();
        assert_eq!(path.states, vec![0, 1, 2]);
        assert_eq!(path.arcs.len(), 2);
        assert_eq!(path.input_labels(), vec![1, 2]);
        assert_eq!(path.output_labels(), vec![1, 2]);
    }

    #[test]
    fn test_decode_linear_fst_no_start() {
        let fst = VectorFst::<TropicalWeight>::new();
        assert!(decode_linear_fst(&fst).is_err());
    }

    #[test]
    fn test_decode_linear_fst_multiple_paths() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::one());
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
        fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s2)); // Multiple paths!
        assert!(decode_linear_fst(&fst).is_err());
    }
}