arcweight 0.3.0

//! Comprehensive integration tests for FailureFst with Aho-Corasick semantics
//!
//! These tests verify correctness of failure transitions, language preservation,
//! and edge cases that could lead to incorrect behavior.

use arcweight::fst::FailureFst;
use arcweight::prelude::*;

#[test]
fn test_failure_fst_creation() {
    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::one(), s1));

    let failure_fst = FailureFst::new(fst);
    assert_eq!(failure_fst.num_states(), 2);
    assert_eq!(failure_fst.start(), Some(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());
    fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

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

    assert_eq!(failure_fst.failure_state(s1), Some(s0));
    assert_eq!(failure_fst.failure_state(s0), None);
}

#[test]
fn test_failure_fst_multiple_failures() {
    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(s1, Arc::new(2, 2, TropicalWeight::one(), s2));

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

    assert_eq!(failure_fst.failure_state(s1), Some(s0));
    assert_eq!(failure_fst.failure_state(s2), Some(s1));
}

#[test]
fn test_failure_fst_arcs() {
    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::one(), s1));

    let failure_fst = FailureFst::new(fst);
    assert_eq!(failure_fst.num_arcs(0), 1);
    assert_eq!(failure_fst.num_arcs(1), 0);
}

#[test]
fn test_failure_fst_final_weight() {
    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!(failure_fst.is_final(1));
    assert_eq!(failure_fst.final_weight(1), Some(&TropicalWeight::new(0.5)));
}

#[test]
fn test_failure_fst_update_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);
    failure_fst.set_failure(s1, s0); // Update to same value

    assert_eq!(failure_fst.failure_state(s1), Some(s0));
}

#[test]
fn test_failure_fst_no_failures() {
    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::one(), s1));

    let failure_fst = FailureFst::new(fst);
    // No failures set, should work like normal FST
    assert_eq!(failure_fst.num_states(), 2);
    assert_eq!(failure_fst.failure_state(0), None);
    assert_eq!(failure_fst.failure_state(1), None);
}

#[test]
fn test_failure_fst_properties() {
    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::one(), s1));

    let failure_fst = FailureFst::new(fst);
    let props = failure_fst.properties();
    // Properties should be preserved - just check it doesn't panic
    let _ = props;
}

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

    for i in 0..100 {
        let s = fst.add_state();
        fst.set_final(s, TropicalWeight::one());
        fst.add_arc(s0, Arc::new(i, i, TropicalWeight::one(), s));
    }

    let mut failure_fst = FailureFst::new(fst);
    // Set failures in a chain
    for i in 1..100 {
        failure_fst.set_failure(i, i - 1);
    }

    assert_eq!(failure_fst.num_states(), 101);
    assert_eq!(failure_fst.failure_state(50), Some(49));
}

/// Test that Aho-Corasick semantics preserve language exactly
#[test]
fn test_aho_corasick_language_preservation() {
    // Create FST: "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::one(), s1),
    );
    fst.add_arc(
        s1,
        Arc::new('b' as u32, 'y' as u32, TropicalWeight::one(), s2),
    );

    let failure_fst = FailureFst::new(fst.clone());

    // Language should be preserved: L(FailureFst) = L(inner FST)
    // Both should accept "ab" -> "xy"
    let paths: Vec<_> = failure_fst.paths_iter().collect();
    let inner_paths: Vec<_> = fst.paths_iter().collect();

    assert_eq!(paths.len(), inner_paths.len());
    assert_eq!(paths[0].input_labels(), inner_paths[0].input_labels());
    assert_eq!(paths[0].output_labels(), inner_paths[0].output_labels());
}

/// Test failure chain with multiple levels
#[test]
fn test_failure_chain_multiple_levels() {
    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();
    fst.set_start(s0);
    fst.set_final(s3, TropicalWeight::one());

    // s0 has arc with 'a'
    fst.add_arc(
        s0,
        Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1),
    );
    // s1 has arc with 'b'
    fst.add_arc(
        s1,
        Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s2),
    );
    // s2 has arc with 'c'
    fst.add_arc(
        s2,
        Arc::new('c' as u32, 'c' as u32, TropicalWeight::one(), s3),
    );

    let mut failure_fst = FailureFst::new(fst);
    // Chain: s3 -> s2 -> s1 -> s0
    failure_fst.set_failure(s3, s2);
    failure_fst.set_failure(s2, s1);
    failure_fst.set_failure(s1, s0);

    // Query 'a' from s3: should find via s2 -> s1 -> s0
    let matching: Vec<_> = failure_fst.arcs_matching(s3, 'a' as u32).collect();
    assert_eq!(matching.len(), 1);
    assert_eq!(matching[0].ilabel, 'a' as u32);
    assert_eq!(matching[0].nextstate, s1);
}

/// Test that multiple arcs with same ilabel in current state are all returned
#[test]
fn test_multiple_arcs_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 but different outputs
    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 mut failure_fst = FailureFst::new(fst);
    failure_fst.set_failure(s1, s0);

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

/// Test epsilon transitions with failure
#[test]
fn test_epsilon_with_failure() {
    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 epsilon arc
    fst.add_arc(s0, Arc::epsilon(TropicalWeight::one(), s1));
    // s1 has arc with 'a'
    fst.add_arc(
        s1,
        Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s2),
    );
    // s0 also has direct arc with 'a' (for failure)
    fst.add_arc(
        s0,
        Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s2),
    );

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

    // Query 'a' from s1: should find arc from s0 via failure
    let matching: Vec<_> = failure_fst.arcs_matching(s1, 'a' as u32).collect();
    assert_eq!(matching.len(), 1);
    assert_eq!(matching[0].ilabel, 'a' as u32);
}

/// Test cycle detection in failure chain
#[test]
fn test_failure_chain_cycle_detection() {
    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::one(), s1));

    let mut failure_fst = FailureFst::new(fst);
    // Create cycle: s0 -> s1 -> s0
    failure_fst.set_failure(s0, s1);
    failure_fst.set_failure(s1, s0);

    // Should not infinite loop - cycle should be detected
    let matching: Vec<_> = failure_fst.arcs_matching(s0, 1).collect();
    // Should find the direct arc from s0
    assert_eq!(matching.len(), 1);
}

/// Test empty state with failure
#[test]
fn test_empty_state_with_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());

    // s0 has arc with 'a'
    fst.add_arc(
        s0,
        Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1),
    );
    // s1 has no arcs

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

    // Query 'a' from s1 (empty state): should find via failure to s0
    let matching: Vec<_> = failure_fst.arcs_matching(s1, 'a' as u32).collect();
    assert_eq!(matching.len(), 1);
    assert_eq!(matching[0].ilabel, 'a' as u32);
}

/// Test that failure is not used when match exists
#[test]
fn test_failure_not_used_when_match_exists() {
    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);

    // Both s0 and s1 have arcs with ilabel=1, but different weights
    fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(10.0), s2));
    fst.add_arc(s1, Arc::new(1, 1, TropicalWeight::new(5.0), s2));

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

    // Query ilabel=1 from s1: should return arc from s1 (weight 5.0), not from s0 (weight 10.0)
    let matching: Vec<_> = failure_fst.arcs_matching(s1, 1).collect();
    assert_eq!(matching.len(), 1);
    assert_eq!(*matching[0].weight.value(), 5.0);
}

/// Test final states with failure transitions
#[test]
fn test_final_states_with_failure() {
    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(s1, TropicalWeight::new(1.0));
    fst.set_final(s2, TropicalWeight::new(2.0));

    fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
    fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));

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

    // Final weights should be preserved
    assert_eq!(
        failure_fst.final_weight(s1),
        Some(&TropicalWeight::new(1.0))
    );
    assert_eq!(
        failure_fst.final_weight(s2),
        Some(&TropicalWeight::new(2.0))
    );
}