arcweight 0.3.0

//! Streaming and real-time optimized algorithms
//!
//! This module provides algorithms optimized for streaming/real-time applications
//! where low latency and bounded memory are important.

use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::{NaturallyOrderedSemiring, Semiring};
use crate::Result;
use std::cmp::Ordering;
use std::collections::BinaryHeap;

/// Streaming shortest path algorithm
///
/// This algorithm processes FSTs in a streaming fashion, producing results
/// incrementally as they become available. Uses Dijkstra's algorithm with
/// a priority queue to ensure shortest paths are found correctly.
#[derive(Debug)]
pub struct StreamingShortestPath<W: Semiring> {
    queue: BinaryHeap<PathState<W>>,
    distances: Vec<Option<W>>,
    /// Maximum number of states to process before stopping
    pub max_states: Option<usize>,
    /// Number of states processed so far
    processed_states: usize,
}

/// State in priority queue for shortest path
#[derive(Clone, Debug)]
struct PathState<W: Semiring> {
    state: StateId,
    weight: W,
}

impl<W: NaturallyOrderedSemiring> PartialEq for PathState<W> {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}

impl<W: NaturallyOrderedSemiring> Eq for PathState<W> {}

impl<W: NaturallyOrderedSemiring> PartialOrd for PathState<W> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<W: NaturallyOrderedSemiring> Ord for PathState<W> {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse ordering for min-heap: smaller weights have higher priority
        other.weight.cmp(&self.weight)
    }
}

impl<W: Semiring> StreamingShortestPath<W>
where
    W: Clone + NaturallyOrderedSemiring,
{
    /// Create a new streaming shortest path processor
    pub fn new(max_states: Option<usize>) -> Self {
        Self {
            queue: BinaryHeap::new(),
            distances: Vec::new(),
            max_states,
            processed_states: 0,
        }
    }

    /// Process next state in streaming fashion
    ///
    /// Returns the next state and its shortest distance from the start state,
    /// or `None` if all states have been processed or the limit is reached.
    pub fn process_next<F: Fst<W>>(&mut self, fst: &F) -> Option<(StateId, W)>
    where
        W: crate::semiring::NaturallyOrderedSemiring,
    {
        let start = fst.start()?;

        // Initialize on first call
        if self.distances.is_empty() {
            self.distances.resize(fst.num_states(), None);
            // Don't set distance here - let the algorithm set it when processing
            // This ensures the state is properly processed through the priority queue
            self.queue.push(PathState {
                state: start,
                weight: W::one(),
            });
        }

        // Check state limit (based on processed states, not queue size)
        if let Some(max) = self.max_states {
            if self.processed_states >= max {
                return None;
            }
        }

        // Process next state from priority queue (Dijkstra's algorithm)
        while let Some(PathState {
            state,
            weight: dist,
        }) = self.queue.pop()
        {
            // Skip if we've already found a better path to this state
            // This can happen because we may add the same state multiple times with different weights
            if let Some(ref best_dist) = self.distances[state as usize] {
                if dist >= *best_dist {
                    continue;
                }
            }

            // Update distance - we've found the shortest path to this state
            self.distances[state as usize] = Some(dist.clone());

            // Process outgoing arcs
            for arc in fst.arcs(state) {
                let new_dist = dist.clone().times(&arc.weight);
                let next_state = arc.nextstate;

                // Check if this is a better path to next_state
                // For naturally ordered semirings, smaller is better (e.g., tropical semiring)
                let should_update = match &self.distances[next_state as usize] {
                    None => true,
                    Some(old_dist) => new_dist < *old_dist,
                };

                if should_update {
                    // Update distance and add to priority queue
                    // Note: We may add the same state multiple times, but the priority queue
                    // ensures we process the best path first
                    self.distances[next_state as usize] = Some(new_dist.clone());
                    self.queue.push(PathState {
                        state: next_state,
                        weight: new_dist,
                    });
                }
            }

            // Increment processed states counter
            self.processed_states += 1;

            // Return this state with its shortest distance
            return Some((state, dist));
        }

        None
    }
}

/// Bounded memory composition
///
/// Performs composition with a memory limit, useful for resource-constrained
/// environments. If the composition would exceed `max_states`, it returns an error
/// immediately when the limit is reached during construction, not after completion.
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::compose_bounded;
///
/// 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::one(), s1));
///
/// 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::one());
/// fst2.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));
///
/// let result = compose_bounded(&fst1, &fst2, 1000)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
pub fn compose_bounded<W, F1, F2>(
    fst1: &F1,
    fst2: &F2,
    max_states: usize,
) -> Result<crate::fst::VectorFst<W>>
where
    W: Semiring + Clone,
    F1: Fst<W>,
    F2: Fst<W>,
{
    use crate::algorithms::compose::ComposeFilter;
    use crate::algorithms::compose::DefaultComposeFilter;
    use std::collections::HashMap;

    let start1 = fst1
        .start()
        .ok_or_else(|| crate::Error::Algorithm("First FST has no start state".into()))?;
    let start2 = fst2
        .start()
        .ok_or_else(|| crate::Error::Algorithm("Second FST has no start state".into()))?;

    let mut result = crate::fst::VectorFst::default();
    let mut state_map = HashMap::new();
    let mut queue = Vec::new();
    let filter = DefaultComposeFilter;

    // create start state
    let start_state = result.add_state();

    // Check limit immediately (after creating start state)
    if result.num_states() > max_states {
        return Err(crate::Error::Algorithm(format!(
            "Composition would exceed limit of {} states (start state alone exceeds limit)",
            max_states
        )));
    }

    result.set_start(start_state);
    state_map.insert(
        (
            start1,
            start2,
            <DefaultComposeFilter as ComposeFilter<W>>::start(),
        ),
        start_state,
    );
    queue.push((
        start1,
        start2,
        <DefaultComposeFilter as ComposeFilter<W>>::start(),
        start_state,
    ));

    // process states
    while let Some((s1, s2, fs, current)) = queue.pop() {
        // handle final states
        if let (Some(w1), Some(w2)) = (fst1.final_weight(s1), fst2.final_weight(s2)) {
            result.set_final(current, w1.times(w2));
        }

        // process arc pairs
        for arc1 in fst1.arcs(s1) {
            for arc2 in fst2.arcs(s2) {
                if let Some((mut arc, next_fs)) = filter.filter_arc(&arc1, &arc2, &fs) {
                    let next_key = (arc1.nextstate, arc2.nextstate, next_fs);

                    let next_state = match state_map.get(&next_key) {
                        Some(&state) => state,
                        None => {
                            let state = result.add_state();

                            // Check limit after adding new state
                            // Use >= for strict enforcement: if we've reached the limit, stop
                            if result.num_states() >= max_states {
                                return Err(crate::Error::Algorithm(format!(
                                    "Composition exceeded limit of {} states (reached {})",
                                    max_states,
                                    result.num_states()
                                )));
                            }

                            state_map.insert(next_key, state);
                            queue.push((arc1.nextstate, arc2.nextstate, next_fs, state));
                            state
                        }
                    };

                    arc.nextstate = next_state;
                    result.add_arc(current, arc);
                }
            }
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_streaming_shortest_path_new() {
        let processor = StreamingShortestPath::<TropicalWeight>::new(None);
        assert!(processor.max_states.is_none());

        let processor2 = StreamingShortestPath::<TropicalWeight>::new(Some(100));
        assert_eq!(processor2.max_states, Some(100));
    }

    #[test]
    fn test_streaming_shortest_path_process_next() {
        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 processor = StreamingShortestPath::<TropicalWeight>::new(None);
        let result = processor.process_next(&fst);
        assert!(result.is_some());
        let (state, weight) = result.unwrap();
        assert_eq!(state, s0);
        assert_eq!(*weight.value(), 0.0);
    }

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

        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::one());
        fst2.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));

        let result = compose_bounded(&fst1, &fst2, 1000).unwrap();
        assert!(result.num_states() > 0);
    }

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

        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::one());
        fst2.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));

        // Should succeed with large limit
        let result = compose_bounded(&fst1, &fst2, 10).unwrap();
        assert!(result.num_states() <= 10);
    }
}