arcweight 0.3.0

A high-performance, modular library for weighted finite state transducers with comprehensive examples and benchmarks
Documentation
//! Tests for FST drawing and visualization

use arcweight::prelude::*;
use arcweight::utils::{draw_fst, draw_fst_default, DrawingConfig};

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

    let dot = draw_fst_default(&fst).unwrap();
    assert!(dot.contains("digraph FST"));
    assert!(dot.contains("rankdir=LR"));
    assert!(dot.contains("0 -> 1"));
}

#[test]
fn test_draw_fst_with_config() {
    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, 2, TropicalWeight::new(1.5), s1));

    let config = DrawingConfig {
        show_weights: true,
        show_state_ids: true,
        show_final_weights: true,
        use_symbols: false,
        horizontal: true,
        node_shape: "box".to_string(),
        node_color: "white".to_string(),
        start_color: "blue".to_string(),
        final_color: "green".to_string(),
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    assert!(dot.contains("digraph FST"));
    assert!(dot.contains("shape=box"));
    assert!(dot.contains("fillcolor=blue")); // start color
    assert!(dot.contains("fillcolor=green")); // final color
    assert!(dot.contains("1:2 / 1.5")); // arc label with weight
}

#[test]
fn test_draw_fst_with_symbols() {
    let mut fst = VectorFst::<TropicalWeight>::new();
    let mut symbols = SymbolTable::new();
    let hello_id = symbols.add_symbol("hello");
    let world_id = symbols.add_symbol("world");

    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(hello_id, world_id, TropicalWeight::one(), s1));

    let config = DrawingConfig {
        use_symbols: true,
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, Some(&symbols), Some(&symbols)).unwrap();
    assert!(dot.contains("hello:world"));
}

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

    let config = DrawingConfig {
        show_weights: false,
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    assert!(!dot.contains("/ 1.5")); // Should not show weight
}

#[test]
fn test_draw_fst_vertical_layout() {
    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 config = DrawingConfig {
        horizontal: false,
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    assert!(!dot.contains("rankdir=LR")); // Should not have LR
}

#[test]
fn test_draw_fst_multiple_states() {
    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 dot = draw_fst_default(&fst).unwrap();
    assert!(dot.contains("0 -> 1"));
    assert!(dot.contains("1 -> 2"));
    assert!(dot.matches("0").count() >= 2); // Should appear multiple times
}

#[test]
fn test_draw_fst_epsilon_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(0, 0, TropicalWeight::one(), s1)); // epsilon

    let dot = draw_fst_default(&fst).unwrap();
    assert!(dot.contains("0 -> 1"));
}

#[test]
fn test_draw_fst_empty() {
    let fst = VectorFst::<TropicalWeight>::new();
    let dot = draw_fst_default(&fst).unwrap();
    assert!(dot.contains("digraph FST"));
    assert!(dot.contains("}"));
}

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

    let dot = draw_fst_default(&fst).unwrap();
    // Should have both start and final styling
    assert!(dot.contains("0"));
}

#[test]
fn test_draw_fst_complex_structure() {
    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(s2, TropicalWeight::new(0.5));
    fst.set_final(s3, TropicalWeight::new(0.3));

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

    let dot = draw_fst_default(&fst).unwrap();
    assert!(dot.contains("0 -> 1"));
    assert!(dot.contains("1 -> 2"));
    assert!(dot.contains("0 -> 3"));
}

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

    let config = DrawingConfig {
        node_shape: "diamond".to_string(),
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    assert!(dot.contains("shape=diamond"));
}

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

    let config = DrawingConfig {
        show_final_weights: true,
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    assert!(dot.contains("0.75")); // Should show final weight
}

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

    let config = DrawingConfig {
        show_final_weights: false,
        ..Default::default()
    };

    let dot = draw_fst(&fst, config, None, None).unwrap();
    // Final weight might still appear in node label, but shouldn't be emphasized
    assert!(dot.contains("digraph FST"));
}