arcweight 0.3.0

//! Lock-free concurrent FST implementation for multi-threaded access.
//!
//! This module provides [`ConcurrentFst`], a thread-safe FST implementation designed for
//! high-performance concurrent read access with support for concurrent modifications.
//! The implementation uses a combination of atomic operations and fine-grained locking
//! to maximize parallelism while ensuring correctness.
//!
//! # Thread Safety Model
//!
//! - **Read operations**: Lock-free for state counts and start state, fine-grained
//!   read locks for arc access
//! - **Write operations**: Atomic updates for scalars, write locks for collections
//! - **Memory reclamation**: Epoch-based pattern for safe deferred cleanup
//!
//! # Architecture
//!
//! ```text
//! ConcurrentFst<W>
//! +--------------------+
//! | start: AtomicU32   |  Lock-free start state access
//! | num_states: Atomic |  Lock-free state count
//! +--------------------+
//! | states: RwLock     |  Coarse lock for state list
//! |   [0..n] --------> |  ConcurrentState<W>
//! +--------------------+  +------------------+
//! | epoch: EpochGuard  |  | final_weight     | RwLock
//! +--------------------+  | arcs: RwLock     | Fine-grained
//!                         | arc_count: Atomic|
//!                         +------------------+
//! ```
//!
//! # Use Cases
//!
//! - **Server applications**: Shared FST state across request handlers
//! - **Parallel decoding**: Concurrent beam search or Viterbi decoding
//! - **Incremental construction**: Building FSTs from parallel data streams
//! - **Read-heavy workloads**: Many readers with occasional updates
//!
//! # Examples
//!
//! ```rust
//! use arcweight::prelude::*;
//! use arcweight::fst::ConcurrentFst;
//! use std::sync::Arc as StdArc;
//! use std::thread;
//!
//! // Create a concurrent FST
//! let cfst = ConcurrentFst::<TropicalWeight>::new();
//! let s0 = cfst.add_state();
//! let s1 = cfst.add_state();
//! cfst.set_start(s0);
//! cfst.set_final(s1, TropicalWeight::one());
//! cfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));
//!
//! // Share across threads
//! let shared = StdArc::new(cfst);
//!
//! let handles: Vec<_> = (0..4).map(|_| {
//!     let fst = StdArc::clone(&shared);
//!     thread::spawn(move || {
//!         // Concurrent read operations
//!         let start = fst.start();
//!         let num_states = fst.num_states();
//!         (start, num_states)
//!     })
//! }).collect();
//!
//! for handle in handles {
//!     let (start, num_states) = handle.join().unwrap();
//!     assert_eq!(start, Some(0));
//!     assert_eq!(num_states, 2);
//! }
//! ```
//!
//! # References
//!
//! - Herlihy, M., & Shavit, N. (2008). *The Art of Multiprocessor Programming*.
//!   Morgan Kaufmann. ISBN: 978-0123705914.
//!
//! - Fraser, K. (2004). Practical Lock-Freedom. Ph.D. Thesis, University of Cambridge.
//!   UCAM-CL-TR-579.

use crate::arc::{Arc, ArcIterator};
use crate::fst::{Fst, MutableFst, StateId, NO_STATE_ID};
use crate::semiring::Semiring;
use crate::Result;

use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
use std::sync::RwLock;

/// Epoch-based memory reclamation for lock-free data structures
///
/// This provides safe memory reclamation for concurrent access patterns
/// using a simplified epoch-based approach.
#[derive(Debug)]
struct EpochGuard {
    /// Global epoch counter
    global_epoch: AtomicU64,
    /// Per-thread epoch tracking (simplified version)
    active_threads: AtomicUsize,
}

impl EpochGuard {
    fn new() -> Self {
        Self {
            global_epoch: AtomicU64::new(0),
            active_threads: AtomicUsize::new(0),
        }
    }

    /// Enter a critical section for reading
    fn pin(&self) -> EpochPin<'_> {
        self.active_threads.fetch_add(1, Ordering::SeqCst);
        EpochPin { guard: self }
    }

    /// Advance the global epoch (called periodically)
    fn advance(&self) {
        self.global_epoch.fetch_add(1, Ordering::SeqCst);
    }

    /// Check if it's safe to reclaim memory from a given epoch
    #[allow(dead_code)]
    fn is_safe_to_reclaim(&self, _epoch: u64) -> bool {
        // Simplified: safe when no threads are active
        self.active_threads.load(Ordering::SeqCst) == 0
    }
}

/// RAII guard for epoch-based critical section
struct EpochPin<'a> {
    guard: &'a EpochGuard,
}

impl Drop for EpochPin<'_> {
    fn drop(&mut self) {
        self.guard.active_threads.fetch_sub(1, Ordering::SeqCst);
    }
}

/// Thread-safe state for concurrent FST
struct ConcurrentState<W: Semiring> {
    /// Final weight (None if not final)
    final_weight: RwLock<Option<W>>,
    /// Outgoing arcs (protected by RwLock for concurrent access)
    arcs: RwLock<Vec<Arc<W>>>,
    /// Arc count (cached for fast access)
    arc_count: AtomicUsize,
}

impl<W: Semiring> ConcurrentState<W> {
    fn new() -> Self {
        Self {
            final_weight: RwLock::new(None),
            arcs: RwLock::new(Vec::new()),
            arc_count: AtomicUsize::new(0),
        }
    }

    fn with_capacity(arc_capacity: usize) -> Self {
        Self {
            final_weight: RwLock::new(None),
            arcs: RwLock::new(Vec::with_capacity(arc_capacity)),
            arc_count: AtomicUsize::new(0),
        }
    }

    fn set_final(&self, weight: W) {
        let mut final_weight = self.final_weight.write().unwrap();
        *final_weight = Some(weight);
    }

    fn delete_final(&self) {
        let mut final_weight = self.final_weight.write().unwrap();
        *final_weight = None;
    }

    fn is_final(&self) -> bool {
        let final_weight = self.final_weight.read().unwrap();
        final_weight.is_some()
    }

    fn get_final_weight(&self) -> Option<W> {
        let final_weight = self.final_weight.read().unwrap();
        final_weight.clone()
    }

    fn add_arc(&self, arc: Arc<W>) {
        let mut arcs = self.arcs.write().unwrap();
        arcs.push(arc);
        self.arc_count.fetch_add(1, Ordering::Release);
    }

    fn num_arcs(&self) -> usize {
        self.arc_count.load(Ordering::Acquire)
    }

    fn get_arcs(&self) -> Vec<Arc<W>> {
        let arcs = self.arcs.read().unwrap();
        arcs.clone()
    }

    fn clear_arcs(&self) {
        let mut arcs = self.arcs.write().unwrap();
        arcs.clear();
        self.arc_count.store(0, Ordering::Release);
    }
}

impl<W: Semiring> Clone for ConcurrentState<W> {
    fn clone(&self) -> Self {
        Self {
            final_weight: RwLock::new(self.final_weight.read().unwrap().clone()),
            arcs: RwLock::new(self.arcs.read().unwrap().clone()),
            arc_count: AtomicUsize::new(self.arc_count.load(Ordering::Acquire)),
        }
    }
}

/// Lock-free concurrent FST implementation for multi-threaded access.
///
/// This FST supports concurrent read and write operations from multiple threads,
/// using a combination of atomic operations and fine-grained reader-writer locks
/// to maximize parallelism while maintaining correctness and linearizability.
///
/// # Type Parameters
///
/// - `W`: The semiring weight type. Must be `Send + Sync` for thread safety.
///
/// # Complexity
///
/// | Operation | Time Complexity | Contention |
/// |-----------|----------------|------------|
/// | `start()` | $`O(1)`$ | Lock-free |
/// | `num_states()` | $`O(1)`$ | Lock-free |
/// | `add_state()` | $`O(1)`$ amortized | Write lock |
/// | `add_arc(s, arc)` | $`O(1)`$ | Read + state lock |
/// | `arcs(s)` | $`O(k)`$ | Read locks |
///
/// # Thread Safety Guarantees
///
/// - **Linearizability**: All operations appear atomic and ordered
/// - **Non-blocking reads**: Read operations never block other reads
/// - **Isolation**: Writes to different states don't block each other
/// - **Consistency**: Single method calls see consistent state views
///
/// # References
///
/// - Herlihy, M., & Shavit, N. (2008). *The Art of Multiprocessor Programming*.
///   Morgan Kaufmann.
pub struct ConcurrentFst<W: Semiring> {
    /// Start state (NO_STATE_ID if not set)
    start: AtomicU32,
    /// State storage
    states: RwLock<Vec<Box<ConcurrentState<W>>>>,
    /// State count (cached)
    num_states: AtomicUsize,
    /// Total arc count (cached)
    total_arcs: AtomicUsize,
    /// Epoch guard for memory safety
    epoch: EpochGuard,
    /// Properties cache
    properties: RwLock<crate::properties::FstProperties>,
}

impl<W: Semiring> ConcurrentFst<W> {
    /// Create a new empty concurrent FST
    pub fn new() -> Self {
        Self {
            start: AtomicU32::new(NO_STATE_ID),
            states: RwLock::new(Vec::new()),
            num_states: AtomicUsize::new(0),
            total_arcs: AtomicUsize::new(0),
            epoch: EpochGuard::new(),
            properties: RwLock::new(crate::properties::FstProperties::default()),
        }
    }

    /// Create a concurrent FST with pre-allocated capacity
    pub fn with_capacity(state_capacity: usize) -> Self {
        Self {
            start: AtomicU32::new(NO_STATE_ID),
            states: RwLock::new(Vec::with_capacity(state_capacity)),
            num_states: AtomicUsize::new(0),
            total_arcs: AtomicUsize::new(0),
            epoch: EpochGuard::new(),
            properties: RwLock::new(crate::properties::FstProperties::default()),
        }
    }

    /// Create a concurrent FST from any FST implementation
    ///
    /// This converts the input FST to a thread-safe concurrent version.
    pub fn from_fst<F: Fst<W>>(fst: &F) -> Result<Self> {
        let num_states = fst.num_states();
        let mut states = Vec::with_capacity(num_states);

        for state_id in 0..num_states as StateId {
            let state = ConcurrentState::with_capacity(fst.num_arcs(state_id));

            // Set final weight
            if let Some(w) = fst.final_weight(state_id) {
                state.set_final(w.clone());
            }

            // Add arcs
            for arc in fst.arcs(state_id) {
                state.add_arc(arc.clone());
            }

            states.push(Box::new(state));
        }

        let total_arcs: usize = states.iter().map(|s| s.num_arcs()).sum();

        Ok(Self {
            start: AtomicU32::new(fst.start().unwrap_or(NO_STATE_ID)),
            states: RwLock::new(states),
            num_states: AtomicUsize::new(num_states),
            total_arcs: AtomicUsize::new(total_arcs),
            epoch: EpochGuard::new(),
            properties: RwLock::new(fst.properties()),
        })
    }

    /// Add a new state and return its ID (thread-safe)
    pub fn add_state(&self) -> StateId {
        let mut states = self.states.write().unwrap();
        let state_id = states.len() as StateId;
        states.push(Box::new(ConcurrentState::new()));
        self.num_states.fetch_add(1, Ordering::Release);
        state_id
    }

    /// Add multiple states at once (more efficient for batch operations)
    pub fn add_states(&self, count: usize) -> Vec<StateId> {
        let mut states = self.states.write().unwrap();
        let start_id = states.len() as StateId;

        for _ in 0..count {
            states.push(Box::new(ConcurrentState::new()));
        }

        self.num_states.fetch_add(count, Ordering::Release);
        (start_id..start_id + count as StateId).collect()
    }

    /// Set the start state (thread-safe)
    pub fn set_start(&self, state: StateId) {
        self.start.store(state, Ordering::Release);
    }

    /// Set the final weight for a state (thread-safe)
    pub fn set_final(&self, state: StateId, weight: W) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            s.set_final(weight);
        }
    }

    /// Delete the final weight for a state (thread-safe)
    pub fn delete_final(&self, state: StateId) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            s.delete_final();
        }
    }

    /// Add an arc to a state (thread-safe)
    pub fn add_arc(&self, state: StateId, arc: Arc<W>) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            s.add_arc(arc);
            self.total_arcs.fetch_add(1, Ordering::Release);
        }
    }

    /// Add multiple arcs at once (more efficient for batch operations)
    pub fn add_arcs(&self, state: StateId, arcs: Vec<Arc<W>>) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            let count = arcs.len();
            for arc in arcs {
                s.add_arc(arc);
            }
            self.total_arcs.fetch_add(count, Ordering::Release);
        }
    }

    /// Clear all arcs from a state (thread-safe)
    pub fn clear_arcs(&self, state: StateId) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            let old_count = s.num_arcs();
            s.clear_arcs();
            self.total_arcs.fetch_sub(old_count, Ordering::Release);
        }
    }

    /// Get the total number of arcs in the FST
    pub fn total_arcs(&self) -> usize {
        self.total_arcs.load(Ordering::Acquire)
    }

    /// Check if a state is final (thread-safe)
    pub fn is_final(&self, state: StateId) -> bool {
        let states = self.states.read().unwrap();
        states
            .get(state as usize)
            .map(|s| s.is_final())
            .unwrap_or(false)
    }

    /// Get final weight (returns owned copy for thread safety)
    pub fn get_final_weight(&self, state: StateId) -> Option<W> {
        let states = self.states.read().unwrap();
        states
            .get(state as usize)
            .and_then(|s| s.get_final_weight())
    }

    /// Get arcs from a state (returns owned copy for thread safety)
    pub fn get_arcs(&self, state: StateId) -> Vec<Arc<W>> {
        let states = self.states.read().unwrap();
        states
            .get(state as usize)
            .map(|s| s.get_arcs())
            .unwrap_or_default()
    }

    /// Execute a function with read access to state data
    ///
    /// This allows efficient access without copying when the callback
    /// processes data immediately.
    #[cfg(test)]
    fn with_state<R, F>(&self, state: StateId, f: F) -> Option<R>
    where
        F: FnOnce(&ConcurrentState<W>) -> R,
    {
        let states = self.states.read().unwrap();
        states.get(state as usize).map(|s| f(s))
    }

    /// Iterate over all states (returns state IDs)
    pub fn states(&self) -> impl Iterator<Item = StateId> {
        let num = self.num_states.load(Ordering::Acquire);
        0..num as StateId
    }

    /// Create a snapshot of the FST for consistent iteration
    ///
    /// This creates a copy of the FST at a point in time, useful when
    /// you need consistent iteration while other threads may be modifying.
    pub fn snapshot(&self) -> ConcurrentFstSnapshot<W> {
        let _pin = self.epoch.pin();
        let states = self.states.read().unwrap();

        let snapshot_states: Vec<SnapshotState<W>> = states
            .iter()
            .map(|s| SnapshotState {
                final_weight: s.get_final_weight(),
                arcs: s.get_arcs(),
            })
            .collect();

        ConcurrentFstSnapshot {
            start: self.start.load(Ordering::Acquire),
            states: snapshot_states,
        }
    }

    /// Compact the FST by removing unreachable states
    ///
    /// Note: This operation temporarily blocks all writers.
    pub fn compact(&self) {
        // Get exclusive access
        let mut states = self.states.write().unwrap();
        let start = self.start.load(Ordering::Acquire);

        if start == NO_STATE_ID || states.is_empty() {
            return;
        }

        // Find reachable states using BFS
        let mut reachable = vec![false; states.len()];
        let mut queue = std::collections::VecDeque::new();

        reachable[start as usize] = true;
        queue.push_back(start);

        while let Some(state) = queue.pop_front() {
            let arcs = states[state as usize].get_arcs();
            for arc in arcs {
                if !reachable[arc.nextstate as usize] {
                    reachable[arc.nextstate as usize] = true;
                    queue.push_back(arc.nextstate);
                }
            }
        }

        // Create mapping from old to new state IDs
        let mut old_to_new: Vec<StateId> = vec![NO_STATE_ID; states.len()];
        let mut new_id = 0;
        for (old_id, &is_reachable) in reachable.iter().enumerate() {
            if is_reachable {
                old_to_new[old_id] = new_id;
                new_id += 1;
            }
        }

        // Create new state list with remapped arc targets
        let mut new_states = Vec::with_capacity(new_id as usize);
        let mut new_total_arcs = 0;

        for (old_id, state) in states.drain(..).enumerate() {
            if reachable[old_id] {
                // Remap arc targets
                let mut arcs = state.get_arcs();
                for arc in &mut arcs {
                    arc.nextstate = old_to_new[arc.nextstate as usize];
                }

                let new_state = ConcurrentState::with_capacity(arcs.len());
                if let Some(w) = state.get_final_weight() {
                    new_state.set_final(w);
                }
                for arc in arcs {
                    new_state.add_arc(arc);
                }
                new_total_arcs += new_state.num_arcs();

                new_states.push(Box::new(new_state));
            }
        }

        // Update state
        *states = new_states;
        self.num_states.store(states.len(), Ordering::Release);
        self.total_arcs.store(new_total_arcs, Ordering::Release);
        self.start
            .store(old_to_new[start as usize], Ordering::Release);

        // Advance epoch to allow memory reclamation
        self.epoch.advance();
    }
}

impl<W: Semiring> Default for ConcurrentFst<W> {
    fn default() -> Self {
        Self::new()
    }
}

impl<W: Semiring> Clone for ConcurrentFst<W> {
    fn clone(&self) -> Self {
        let states = self.states.read().unwrap();
        let cloned_states: Vec<Box<ConcurrentState<W>>> =
            states.iter().map(|s| Box::new((**s).clone())).collect();

        Self {
            start: AtomicU32::new(self.start.load(Ordering::Acquire)),
            states: RwLock::new(cloned_states),
            num_states: AtomicUsize::new(self.num_states.load(Ordering::Acquire)),
            total_arcs: AtomicUsize::new(self.total_arcs.load(Ordering::Acquire)),
            epoch: EpochGuard::new(),
            properties: RwLock::new(*self.properties.read().unwrap()),
        }
    }
}

impl<W: Semiring> std::fmt::Debug for ConcurrentFst<W> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConcurrentFst")
            .field("num_states", &self.num_states.load(Ordering::Acquire))
            .field("total_arcs", &self.total_arcs.load(Ordering::Acquire))
            .field("start", &self.start.load(Ordering::Acquire))
            .finish()
    }
}

// Thread-safe trait implementations
unsafe impl<W: Semiring + Send> Send for ConcurrentFst<W> {}
unsafe impl<W: Semiring + Send + Sync> Sync for ConcurrentFst<W> {}

/// Iterator over arcs from a concurrent FST state
#[derive(Debug)]
pub struct ConcurrentArcIterator<W: Semiring> {
    arcs: Vec<Arc<W>>,
    index: usize,
}

impl<W: Semiring> Iterator for ConcurrentArcIterator<W> {
    type Item = Arc<W>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < self.arcs.len() {
            let arc = self.arcs[self.index].clone();
            self.index += 1;
            Some(arc)
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.arcs.len() - self.index;
        (remaining, Some(remaining))
    }
}

impl<W: Semiring> ExactSizeIterator for ConcurrentArcIterator<W> {}

impl<W: Semiring> ArcIterator<W> for ConcurrentArcIterator<W> {
    fn reset(&mut self) {
        self.index = 0;
    }
}

impl<W: Semiring> Fst<W> for ConcurrentFst<W> {
    type ArcIter<'a>
        = ConcurrentArcIterator<W>
    where
        W: 'a;

    fn start(&self) -> Option<StateId> {
        let start = self.start.load(Ordering::Acquire);
        if start == NO_STATE_ID {
            None
        } else {
            Some(start)
        }
    }

    fn final_weight(&self, _state: StateId) -> Option<&W> {
        // Note: This returns None because we can't return a reference to
        // lock-protected data. Use get_final_weight() for thread-safe access.
        // This is a limitation of the Fst trait design.
        None
    }

    fn num_states(&self) -> usize {
        self.num_states.load(Ordering::Acquire)
    }

    fn num_arcs(&self, state: StateId) -> usize {
        let states = self.states.read().unwrap();
        states
            .get(state as usize)
            .map(|s| s.num_arcs())
            .unwrap_or(0)
    }

    fn arcs(&self, state: StateId) -> Self::ArcIter<'_> {
        let arcs = self.get_arcs(state);
        ConcurrentArcIterator { arcs, index: 0 }
    }

    fn properties(&self) -> crate::properties::FstProperties {
        *self.properties.read().unwrap()
    }
}

impl<W: Semiring> MutableFst<W> for ConcurrentFst<W> {
    fn add_state(&mut self) -> StateId {
        ConcurrentFst::add_state(self)
    }

    fn set_start(&mut self, state: StateId) {
        ConcurrentFst::set_start(self, state)
    }

    fn set_final(&mut self, state: StateId, weight: W) {
        ConcurrentFst::set_final(self, state, weight)
    }

    fn add_arc(&mut self, state: StateId, arc: Arc<W>) {
        ConcurrentFst::add_arc(self, state, arc)
    }

    fn delete_arcs(&mut self, state: StateId) {
        ConcurrentFst::clear_arcs(self, state)
    }

    fn delete_arc(&mut self, state: StateId, arc_idx: usize) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            let mut arcs = s.arcs.write().unwrap();
            if arc_idx < arcs.len() {
                arcs.remove(arc_idx);
                s.arc_count.fetch_sub(1, Ordering::Release);
                self.total_arcs.fetch_sub(1, Ordering::Release);
            }
        }
    }

    fn reserve_states(&mut self, n: usize) {
        let mut states = self.states.write().unwrap();
        states.reserve(n);
    }

    fn reserve_arcs(&mut self, state: StateId, n: usize) {
        let states = self.states.read().unwrap();
        if let Some(s) = states.get(state as usize) {
            let mut arcs = s.arcs.write().unwrap();
            arcs.reserve(n);
        }
    }

    fn clear(&mut self) {
        let mut states = self.states.write().unwrap();
        states.clear();
        self.start.store(NO_STATE_ID, Ordering::Release);
        self.num_states.store(0, Ordering::Release);
        self.total_arcs.store(0, Ordering::Release);
    }
}

/// Snapshot state for consistent iteration
#[derive(Debug)]
struct SnapshotState<W: Semiring> {
    final_weight: Option<W>,
    arcs: Vec<Arc<W>>,
}

/// Immutable snapshot of a ConcurrentFst
///
/// This provides a consistent view of the FST at a point in time,
/// allowing iteration without interference from concurrent modifications.
#[derive(Debug)]
pub struct ConcurrentFstSnapshot<W: Semiring> {
    start: StateId,
    states: Vec<SnapshotState<W>>,
}

impl<W: Semiring> ConcurrentFstSnapshot<W> {
    /// Get the start state
    pub fn start(&self) -> Option<StateId> {
        if self.start == NO_STATE_ID {
            None
        } else {
            Some(self.start)
        }
    }

    /// Get the number of states
    pub fn num_states(&self) -> usize {
        self.states.len()
    }

    /// Get final weight for a state
    pub fn final_weight(&self, state: StateId) -> Option<&W> {
        self.states
            .get(state as usize)
            .and_then(|s| s.final_weight.as_ref())
    }

    /// Get arcs for a state
    pub fn arcs(&self, state: StateId) -> impl Iterator<Item = &Arc<W>> {
        self.states
            .get(state as usize)
            .map(|s| s.arcs.iter())
            .into_iter()
            .flatten()
    }

    /// Check if a state is final
    pub fn is_final(&self, state: StateId) -> bool {
        self.states
            .get(state as usize)
            .map(|s| s.final_weight.is_some())
            .unwrap_or(false)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::prelude::*;
    use std::sync::Arc as StdArc;
    use std::thread;

    #[test]
    fn test_concurrent_fst_basic() {
        let fst = ConcurrentFst::<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, 2, TropicalWeight::new(1.0), s1));

        assert_eq!(fst.start(), Some(0));
        assert_eq!(fst.num_states(), 2);
        assert_eq!(fst.num_arcs(s0), 1);
        assert!(fst.is_final(s1));
    }

    #[test]
    fn test_concurrent_fst_from_fst() {
        let mut vector_fst = VectorFst::<TropicalWeight>::new();
        let s0 = vector_fst.add_state();
        let s1 = vector_fst.add_state();
        vector_fst.set_start(s0);
        vector_fst.set_final(s1, TropicalWeight::one());
        vector_fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));

        let cfst = ConcurrentFst::from_fst(&vector_fst).unwrap();

        assert_eq!(cfst.start(), Some(0));
        assert_eq!(cfst.num_states(), 2);
        assert_eq!(cfst.num_arcs(s0), 1);
        assert_eq!(cfst.get_final_weight(s1), Some(TropicalWeight::one()));
    }

    #[test]
    fn test_concurrent_fst_multithreaded_read() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        cfst.set_start(s0);
        cfst.set_final(s1, TropicalWeight::one());

        for i in 1..=10 {
            cfst.add_arc(s0, Arc::new(i, i, TropicalWeight::new(i as f32), s1));
        }

        let shared = StdArc::new(cfst);

        // Spawn multiple reader threads
        let handles: Vec<_> = (0..8)
            .map(|_| {
                let fst = StdArc::clone(&shared);
                thread::spawn(move || {
                    for _ in 0..100 {
                        let start = fst.start();
                        let num_states = fst.num_states();
                        let num_arcs = fst.num_arcs(0);
                        let arcs = fst.get_arcs(0);

                        assert_eq!(start, Some(0));
                        assert_eq!(num_states, 2);
                        assert_eq!(num_arcs, 10);
                        assert_eq!(arcs.len(), 10);
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
    }

    #[test]
    fn test_concurrent_fst_snapshot() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        cfst.set_start(s0);
        cfst.set_final(s1, TropicalWeight::new(0.5));
        cfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));

        let snapshot = cfst.snapshot();

        assert_eq!(snapshot.start(), Some(0));
        assert_eq!(snapshot.num_states(), 2);
        assert!(snapshot.is_final(s1));

        let arcs: Vec<_> = snapshot.arcs(s0).collect();
        assert_eq!(arcs.len(), 1);
        assert_eq!(arcs[0].ilabel, 1);
    }

    #[test]
    fn test_concurrent_fst_batch_operations() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();

        // Add states in batch
        let states = cfst.add_states(100);
        assert_eq!(states.len(), 100);
        assert_eq!(cfst.num_states(), 100);

        // Add arcs in batch
        let arcs: Vec<Arc<TropicalWeight>> = (1..=10)
            .map(|i| Arc::new(i, i, TropicalWeight::new(i as f32), 1))
            .collect();
        cfst.add_arcs(0, arcs);

        assert_eq!(cfst.num_arcs(0), 10);
        assert_eq!(cfst.total_arcs(), 10);
    }

    #[test]
    fn test_concurrent_fst_clear_arcs() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();

        for i in 1..=5 {
            cfst.add_arc(s0, Arc::new(i, i, TropicalWeight::new(i as f32), s1));
        }

        assert_eq!(cfst.num_arcs(s0), 5);
        assert_eq!(cfst.total_arcs(), 5);

        cfst.clear_arcs(s0);

        assert_eq!(cfst.num_arcs(s0), 0);
        assert_eq!(cfst.total_arcs(), 0);
    }

    #[test]
    fn test_concurrent_fst_compact() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();

        // Create states
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        let s2 = cfst.add_state(); // unreachable
        let s3 = cfst.add_state();

        cfst.set_start(s0);
        cfst.set_final(s3, TropicalWeight::one());

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

        // s2 is unreachable
        cfst.set_final(s2, TropicalWeight::one());

        assert_eq!(cfst.num_states(), 4);

        cfst.compact();

        // Should have removed unreachable s2
        assert_eq!(cfst.num_states(), 3);
    }

    #[test]
    fn test_concurrent_fst_clone() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        cfst.set_start(s0);
        cfst.set_final(s1, TropicalWeight::one());
        cfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));

        let cloned = cfst.clone();

        assert_eq!(cloned.start(), Some(0));
        assert_eq!(cloned.num_states(), 2);
        assert_eq!(cloned.num_arcs(s0), 1);

        // Modifications to clone don't affect original
        cloned.add_arc(s0, Arc::new(2, 2, TropicalWeight::new(1.0), s1));
        assert_eq!(cfst.num_arcs(s0), 1);
        assert_eq!(cloned.num_arcs(s0), 2);
    }

    #[test]
    fn test_concurrent_fst_arc_iterator() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        cfst.set_start(s0);
        cfst.set_final(s1, TropicalWeight::one());

        for i in 1..=5 {
            cfst.add_arc(s0, Arc::new(i, i * 10, TropicalWeight::new(i as f32), s1));
        }

        let arcs: Vec<_> = cfst.arcs(s0).collect();
        assert_eq!(arcs.len(), 5);

        for (i, arc) in arcs.iter().enumerate() {
            assert_eq!(arc.ilabel, (i + 1) as u32);
            assert_eq!(arc.olabel, ((i + 1) * 10) as u32);
        }
    }

    #[test]
    fn test_concurrent_fst_with_state() {
        let cfst = ConcurrentFst::<TropicalWeight>::new();
        let s0 = cfst.add_state();
        let s1 = cfst.add_state();
        cfst.set_start(s0);
        cfst.set_final(s1, TropicalWeight::new(0.5));
        cfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

        // Use with_state for efficient access
        let result = cfst.with_state(s0, |state| (state.num_arcs(), state.is_final()));

        assert_eq!(result, Some((1, false)));

        let result2 = cfst.with_state(s1, |state| (state.num_arcs(), state.is_final()));

        assert_eq!(result2, Some((0, true)));
    }
}