arcweight 0.3.0

//! Failure transitions for compact FST representation using Aho-Corasick semantics.
//!
//! This module provides [`FailureFst`], a wrapper that adds failure (fallback) transitions
//! to any FST implementation. Failure transitions enable compact representation of large
//! automata by sharing common suffix states, following the pattern established by the
//! Aho-Corasick string matching algorithm.
//!
//! # Failure Transition Semantics
//!
//! This implementation uses true Aho-Corasick semantics where failure transitions are
//! only traversed when no matching arc exists for a given input label. This preserves
//! the exact language of the underlying FST: $`L(\text{FailureFst}) = L(\text{inner FST})`$.
//!
//! # Algorithm
//!
//! For state $`s`$ and input label $`\sigma`$:
//!
//! 1. If $`s`$ has an arc with input label $`\sigma`$, use that arc
//! 2. Otherwise, let $`s' = \text{failure}(s)`$ and repeat from step 1 with $`s'`$
//! 3. If failure chain is exhausted with no match, no transition exists
//!
//! This conditional semantics ensures deterministic matching while allowing significant
//! state sharing for prefix/suffix-closed languages.
//!
//! # Complexity
//!
//! - **Space**: $`O(V)`$ additional storage for failure map
//! - **Arc lookup with `arcs_matching`**: $`O(d \cdot h)`$ where $`d`$ is average out-degree
//!   and $`h`$ is failure chain height
//! - **Arc iteration with `arcs`**: Returns all arcs from state plus failure chain
//!
//! # Examples
//!
//! ```rust
//! use arcweight::prelude::*;
//! use arcweight::fst::FailureFst;
//!
//! // Create base FST
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! let s0 = fst.add_state();
//! let s1 = fst.add_state();
//! fst.set_start(s0);
//! fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
//!
//! // Wrap with failure transitions
//! let mut failure_fst = FailureFst::new(fst);
//! failure_fst.set_failure(s1, s0); // s1 fails back to s0
//!
//! // Query arcs matching specific label (uses failure semantics)
//! let matching: Vec<_> = failure_fst.arcs_matching(s1, 1).collect();
//! ```
//!
//! # References
//!
//! - Aho, A. V., & Corasick, M. J. (1975). Efficient String Matching: An Aid to
//!   Bibliographic Search. *Communications of the ACM*, 18(6), 333-340.
//!
//! - Mohri, M. (1997). Finite-State Transducers in Language and Speech Processing.
//!   *Computational Linguistics*, 23(2), 269-311.
//!
//! - Allauzen, C., & Mohri, M. (2006). Efficient Algorithms for Testing the Twins
//!   Property. *Journal of Automata, Languages and Combinatorics*, 8(2), 117-144.

use crate::arc::{Arc, ArcIterator};
use crate::fst::{Fst, Label, StateId};
use crate::semiring::Semiring;
use std::collections::{HashMap, HashSet};

/// FST wrapper with Aho-Corasick style failure transitions.
///
/// This wrapper adds failure transition support to an existing FST, enabling
/// compact representation of large automata through state sharing. Failure
/// transitions are followed only when no matching arc exists at the current
/// state, preserving the language of the underlying FST.
///
/// # Type Parameters
///
/// - `F`: The underlying FST type
/// - `W`: The semiring weight type
///
/// # References
///
/// - Aho, A. V., & Corasick, M. J. (1975). Efficient String Matching: An Aid to
///   Bibliographic Search. *Communications of the ACM*, 18(6), 333-340.
#[derive(Debug)]
pub struct FailureFst<F: Fst<W>, W: Semiring> {
    inner: F,
    failure_map: HashMap<StateId, StateId>, // state -> failure_state
    _phantom: std::marker::PhantomData<W>,
}

/// Iterator over arcs from a state, including arcs from failure states
///
/// This iterator implements true Aho-Corasick semantics: it returns all arcs
/// from the current state, then arcs from failure states (for backward compatibility
/// with the `arcs()` method). For conditional failure semantics, use `arcs_matching()`.
pub struct FailureArcIterator<'a, W: Semiring, F: Fst<W>> {
    inner: &'a F,
    _failure_map: &'a HashMap<StateId, StateId>, // Used in new() to build failure_chain
    current_state: StateId,
    inner_arcs: <F as Fst<W>>::ArcIter<'a>,
    failure_chain: Vec<StateId>, // States to visit via failure transitions
    _visited_failures: HashSet<StateId>, // Used in new() to detect cycles
    failure_arcs: Option<<F as Fst<W>>::ArcIter<'a>>,
    failure_state_idx: usize,
}

/// Iterator over arcs matching a specific input label with Aho-Corasick semantics
///
/// Returns arcs from the current state first. If no matching arc is found,
/// follows failure transitions until a match is found or the failure chain is exhausted.
pub struct FailureMatchingArcIterator<'a, W: Semiring, F: Fst<W>> {
    inner: &'a F,
    #[allow(dead_code)] // Used indirectly through failure_chain built in new()
    failure_map: &'a HashMap<StateId, StateId>,
    current_state: StateId,
    target_ilabel: Label,
    inner_arcs: <F as Fst<W>>::ArcIter<'a>,
    failure_chain: Vec<StateId>,
    failure_state_idx: usize,
    found_match: bool,
}

impl<'a, W: Semiring, F: Fst<W>> std::fmt::Debug for FailureMatchingArcIterator<'a, W, F> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FailureMatchingArcIterator")
            .field("current_state", &self.current_state)
            .field("target_ilabel", &self.target_ilabel)
            .field("failure_chain", &self.failure_chain)
            .field("failure_state_idx", &self.failure_state_idx)
            .field("found_match", &self.found_match)
            .finish_non_exhaustive()
    }
}

impl<'a, W: Semiring, F: Fst<W>> std::fmt::Debug for FailureArcIterator<'a, W, F> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FailureArcIterator")
            .field("current_state", &self.current_state)
            .field("failure_chain", &self.failure_chain)
            .field("failure_state_idx", &self.failure_state_idx)
            .finish_non_exhaustive()
    }
}

impl<'a, W: Semiring, F: Fst<W>> FailureArcIterator<'a, W, F> {
    fn new(inner: &'a F, failure_map: &'a HashMap<StateId, StateId>, state: StateId) -> Self {
        let inner_arcs = inner.arcs(state);
        let mut failure_chain = Vec::new();
        let mut visited = HashSet::new();

        // Build failure chain starting from this state
        let mut current = state;
        while let Some(&failure_state) = failure_map.get(&current) {
            if visited.contains(&failure_state) {
                break; // Cycle detected
            }
            visited.insert(failure_state);
            failure_chain.push(failure_state);
            current = failure_state;
        }

        Self {
            inner,
            _failure_map: failure_map,
            current_state: state,
            inner_arcs,
            failure_chain,
            _visited_failures: visited,
            failure_arcs: None,
            failure_state_idx: 0,
        }
    }
}

impl<'a, W: Semiring, F: Fst<W>> Iterator for FailureArcIterator<'a, W, F> {
    type Item = Arc<W>;

    fn next(&mut self) -> Option<Self::Item> {
        // First, yield all arcs from the inner FST
        if let Some(arc) = self.inner_arcs.next() {
            return Some(arc);
        }

        // Then, yield arcs from failure states
        while self.failure_state_idx < self.failure_chain.len() {
            let failure_state = self.failure_chain[self.failure_state_idx];

            // Initialize iterator for this failure state if needed
            if self.failure_arcs.is_none() {
                self.failure_arcs = Some(self.inner.arcs(failure_state));
            }

            // Get next arc from current failure state
            if let Some(ref mut iter) = self.failure_arcs {
                if let Some(arc) = iter.next() {
                    return Some(arc);
                }
            }

            // Move to next failure state
            self.failure_arcs = None;
            self.failure_state_idx += 1;
        }

        None
    }
}

impl<'a, W: Semiring, F: Fst<W>> ArcIterator<W> for FailureArcIterator<'a, W, F> {
    fn reset(&mut self) {
        // Reset inner arcs iterator
        self.inner_arcs = self.inner.arcs(self.current_state);

        // Reset failure state iteration
        self.failure_arcs = None;
        self.failure_state_idx = 0;
    }
}

impl<'a, W: Semiring, F: Fst<W>> FailureMatchingArcIterator<'a, W, F> {
    fn new(
        inner: &'a F,
        failure_map: &'a HashMap<StateId, StateId>,
        state: StateId,
        target_ilabel: Label,
    ) -> Self {
        let inner_arcs = inner.arcs(state);
        let mut failure_chain = Vec::new();
        let mut visited = HashSet::new();

        // Build failure chain starting from this state
        let mut current = state;
        while let Some(&failure_state) = failure_map.get(&current) {
            if visited.contains(&failure_state) {
                break; // Cycle detected
            }
            visited.insert(failure_state);
            failure_chain.push(failure_state);
            current = failure_state;
        }

        Self {
            inner,
            failure_map,
            current_state: state,
            target_ilabel,
            inner_arcs,
            failure_chain,
            failure_state_idx: 0,
            found_match: false,
        }
    }
}

impl<'a, W: Semiring, F: Fst<W>> Iterator for FailureMatchingArcIterator<'a, W, F> {
    type Item = Arc<W>;

    fn next(&mut self) -> Option<Self::Item> {
        // First, check arcs from current state
        for arc in self.inner_arcs.by_ref() {
            if arc.ilabel == self.target_ilabel {
                self.found_match = true;
                return Some(arc);
            }
        }

        // If we found a match in current state, we're done (Aho-Corasick semantics)
        if self.found_match {
            return None;
        }

        // No match in current state - follow failure chain
        while self.failure_state_idx < self.failure_chain.len() {
            let failure_state = self.failure_chain[self.failure_state_idx];
            let failure_arcs = self.inner.arcs(failure_state);

            // Look for matching arc in this failure state
            for arc in failure_arcs {
                if arc.ilabel == self.target_ilabel {
                    self.found_match = true;
                    return Some(arc);
                }
            }

            // No match in this failure state, try next
            self.failure_state_idx += 1;
        }

        None
    }
}

impl<F: Fst<W>, W: Semiring> FailureFst<F, W> {
    /// Create a new failure FST from an existing FST
    pub fn new(fst: F) -> Self {
        Self {
            inner: fst,
            failure_map: HashMap::new(),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Set a failure transition
    pub fn set_failure(&mut self, state: StateId, failure_state: StateId) {
        self.failure_map.insert(state, failure_state);
    }

    /// Get the failure state for a given state
    pub fn failure_state(&self, state: StateId) -> Option<StateId> {
        self.failure_map.get(&state).copied()
    }

    /// Get arcs matching a specific input label with Aho-Corasick semantics
    ///
    /// This method implements true Aho-Corasick conditional failure transitions:
    /// 1. First checks for arcs in the current state matching `ilabel`
    /// 2. If no match found, follows failure chain until a match is found
    /// 3. Returns only matching arcs, preserving exact language semantics
    ///
    /// # Examples
    ///
    /// ```rust
    /// use arcweight::prelude::*;
    /// use arcweight::fst::FailureFst;
    ///
    /// let mut fst = VectorFst::<TropicalWeight>::new();
    /// let s0 = fst.add_state();
    /// let s1 = fst.add_state();
    /// fst.set_start(s0);
    /// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
    ///
    /// let mut failure_fst = FailureFst::new(fst);
    /// failure_fst.set_failure(s1, s0);
    ///
    /// // Get arcs matching input label 1 from state s1
    /// let matching: Vec<_> = failure_fst.arcs_matching(s1, 1).collect();
    /// // Will find arc from s0 (via failure transition) if s1 has no arc with ilabel=1
    /// ```
    pub fn arcs_matching(
        &self,
        state: StateId,
        ilabel: Label,
    ) -> FailureMatchingArcIterator<'_, W, F> {
        FailureMatchingArcIterator::new(&self.inner, &self.failure_map, state, ilabel)
    }
}

impl<F: Fst<W>, W: Semiring> Fst<W> for FailureFst<F, W> {
    type ArcIter<'a>
        = FailureArcIterator<'a, W, F>
    where
        Self: 'a;

    fn start(&self) -> Option<StateId> {
        self.inner.start()
    }

    fn final_weight(&self, state: StateId) -> Option<&W> {
        self.inner.final_weight(state)
    }

    fn num_arcs(&self, state: StateId) -> usize {
        // Count includes arcs from failure states
        // This is an approximation - actual count may vary due to deduplication
        let mut count = self.inner.num_arcs(state);
        let mut visited = HashSet::new();
        let mut current = state;

        while let Some(&failure_state) = self.failure_map.get(&current) {
            if visited.contains(&failure_state) {
                break; // Cycle detected
            }
            visited.insert(failure_state);
            count += self.inner.num_arcs(failure_state);
            current = failure_state;
        }

        count
    }

    fn num_states(&self) -> usize {
        self.inner.num_states()
    }

    fn properties(&self) -> crate::properties::FstProperties {
        self.inner.properties()
    }

    fn arcs(&self, state: StateId) -> Self::ArcIter<'_> {
        FailureArcIterator::new(&self.inner, &self.failure_map, state)
    }
}

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

    #[test]
    fn test_failure_fst_new() {
        let fst = VectorFst::<TropicalWeight>::new();
        let failure_fst = FailureFst::new(fst);
        assert_eq!(failure_fst.num_states(), 0);
    }

    #[test]
    fn test_failure_fst_set_failure() {
        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());

        let mut failure_fst = FailureFst::new(fst);
        failure_fst.set_failure(s1, s0);
        assert_eq!(failure_fst.failure_state(s1), Some(s0));
    }

    #[test]
    fn test_failure_fst_delegates_to_inner() {
        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::new(0.5));
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

        let failure_fst = FailureFst::new(fst);
        assert_eq!(failure_fst.start(), Some(s0));
        assert!(failure_fst.is_final(s1));
        assert_eq!(failure_fst.num_arcs(s0), 1);
    }

    #[test]
    fn test_aho_corasick_conditional_failure() {
        // Test true Aho-Corasick semantics: failure transitions only used when no matching arc
        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());

        // s0 has arc with ilabel=1
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
        // s1 has arc with ilabel=2
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));

        let mut failure_fst = FailureFst::new(fst);
        // s1's failure points to s0
        failure_fst.set_failure(s1, s0);

        // Query arcs matching ilabel=1 from s1
        // Should find arc from s0 (via failure) since s1 has no arc with ilabel=1
        let matching: Vec<_> = failure_fst.arcs_matching(s1, 1).collect();
        assert_eq!(matching.len(), 1);
        assert_eq!(matching[0].ilabel, 1);
        assert_eq!(matching[0].nextstate, s1);

        // Query arcs matching ilabel=2 from s1
        // Should find arc from s1 directly (no failure needed)
        let matching: Vec<_> = failure_fst.arcs_matching(s1, 2).collect();
        assert_eq!(matching.len(), 1);
        assert_eq!(matching[0].ilabel, 2);
        assert_eq!(matching[0].nextstate, s2);
    }

    #[test]
    fn test_aho_corasick_no_failure_when_match_exists() {
        // Test that failure transitions are NOT used when a matching arc exists
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        fst.set_start(s0);

        // Both s0 and s1 have arcs with ilabel=1
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(1, 1, TropicalWeight::new(2.0), s0));

        let mut failure_fst = FailureFst::new(fst);
        failure_fst.set_failure(s1, s0);

        // Query arcs matching ilabel=1 from s1
        // Should only return arc from s1 (not from s0 via failure)
        let matching: Vec<_> = failure_fst.arcs_matching(s1, 1).collect();
        assert_eq!(matching.len(), 1);
        assert_eq!(matching[0].ilabel, 1);
        assert_eq!(*matching[0].weight.value(), 2.0); // Weight from s1, not s0
    }

    #[test]
    fn test_aho_corasick_no_match_found() {
        // Test when no matching arc exists in current state or failure chain
        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());

        // s0 has arc with ilabel=1, s1 has no arcs
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

        let mut failure_fst = FailureFst::new(fst);
        failure_fst.set_failure(s1, s0);

        // Query ilabel=999 from s1: should find nothing (s1 has no arc, s0 has no matching arc)
        let matching: Vec<_> = failure_fst.arcs_matching(s1, 999).collect();
        assert_eq!(matching.len(), 0);
    }

    #[test]
    fn test_aho_corasick_multiple_matches_in_current_state() {
        // Test when current state has multiple arcs with same ilabel
        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);

        // s0 has two arcs with ilabel=1
        fst.add_arc(s0, Arc::new(1, 10, TropicalWeight::new(1.0), s1));
        fst.add_arc(s0, Arc::new(1, 20, TropicalWeight::new(2.0), s2));

        let failure_fst = FailureFst::new(fst);

        // Query ilabel=1 from s0: should return both arcs
        let matching: Vec<_> = failure_fst.arcs_matching(s0, 1).collect();
        assert_eq!(matching.len(), 2);

        // Verify both arcs are present
        let outputs: Vec<u32> = matching.iter().map(|a| a.olabel).collect();
        assert!(outputs.contains(&10));
        assert!(outputs.contains(&20));
    }
}