arcbox-vmm 0.1.4

//! Interrupt controller management.
//!
//! This module provides the IRQ chip abstraction for managing interrupts,
//! including GSI mapping, trigger modes, and interrupt coalescing.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};

use crate::error::Result;

/// IRQ number type.
pub type Irq = u32;

/// Global System Interrupt number type.
pub type Gsi = u32;

/// Maximum number of IRQs.
pub const MAX_IRQS: u32 = 256;

/// Maximum number of GSIs (typically matches IOAPIC entries + legacy PICs).
pub const MAX_GSIS: u32 = 24;

/// Interrupt trigger mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TriggerMode {
    /// Edge-triggered: interrupt is signaled on level transition.
    /// Device asserts then deasserts the IRQ line.
    #[default]
    Edge,
    /// Level-triggered: interrupt remains asserted until acknowledged.
    /// Device keeps line asserted until serviced.
    Level,
}

/// IRQ configuration for a single interrupt line.
#[derive(Debug, Clone)]
pub struct IrqConfig {
    /// The GSI this IRQ is mapped to.
    pub gsi: Gsi,
    /// Trigger mode for this IRQ.
    pub trigger_mode: TriggerMode,
    /// Whether this IRQ is currently asserted (for level-triggered).
    pub asserted: bool,
}

/// Callback type for triggering interrupts on the hypervisor.
///
/// The callback receives (gsi, level) where level is true for assert.
pub type IrqTriggerCallback = Box<dyn Fn(Gsi, bool) -> Result<()> + Send + Sync>;

/// Statistics for interrupt coalescing.
#[derive(Debug, Default)]
pub struct IrqStats {
    /// Total interrupts triggered.
    pub triggered: AtomicU64,
    /// Interrupts coalesced (not delivered because pending).
    pub coalesced: AtomicU64,
}

/// IRQ chip abstraction.
///
/// Manages interrupt routing, delivery, and coalescing.
pub struct IrqChip {
    /// Next available IRQ number.
    next_irq: AtomicU32,
    /// IRQ mask (bit set = masked).
    mask: AtomicU32,
    /// IRQ to GSI mapping and configuration.
    irq_configs: RwLock<HashMap<Irq, IrqConfig>>,
    /// Pending interrupts bitmap for coalescing.
    /// If an IRQ is already pending, we don't trigger again.
    pending: AtomicU32,
    /// Callback for actually triggering interrupts on the VM.
    trigger_callback: Mutex<Option<Arc<IrqTriggerCallback>>>,
    /// Statistics for monitoring.
    stats: IrqStats,
}

impl IrqChip {
    /// Creates a new IRQ chip.
    ///
    /// # Errors
    ///
    /// Returns an error if the IRQ chip cannot be created.
    pub fn new() -> Result<Self> {
        tracing::debug!("Creating IRQ chip");
        Ok(Self {
            next_irq: AtomicU32::new(32), // Start after legacy IRQs
            mask: AtomicU32::new(0),
            irq_configs: RwLock::new(HashMap::new()),
            pending: AtomicU32::new(0),
            trigger_callback: Mutex::new(None),
            stats: IrqStats::default(),
        })
    }

    /// Sets the callback for triggering interrupts.
    ///
    /// This should be called after the VM is created with a callback that
    /// invokes the hypervisor's interrupt injection mechanism.
    pub fn set_trigger_callback(&self, callback: Arc<IrqTriggerCallback>) {
        let mut cb = self
            .trigger_callback
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        *cb = Some(callback);
        tracing::debug!("IRQ trigger callback registered");
    }

    /// Allocates an IRQ number with the specified configuration.
    ///
    /// # Arguments
    /// * `gsi` - The GSI to map this IRQ to
    /// * `trigger_mode` - Edge or level triggered
    ///
    /// # Errors
    ///
    /// Returns an error if no IRQ is available.
    pub fn allocate_irq_with_config(&self, gsi: Gsi, trigger_mode: TriggerMode) -> Result<Irq> {
        let irq = self.next_irq.fetch_add(1, Ordering::SeqCst);
        if irq >= MAX_IRQS {
            return Err(crate::error::VmmError::Irq("IRQ exhausted".to_string()));
        }

        let config = IrqConfig {
            gsi,
            trigger_mode,
            asserted: false,
        };

        {
            let mut configs = self
                .irq_configs
                .write()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            configs.insert(irq, config);
        }

        tracing::debug!(
            "Allocated IRQ {} -> GSI {}, mode={:?}",
            irq,
            gsi,
            trigger_mode
        );

        Ok(irq)
    }

    /// Allocates an IRQ number with default edge-triggered mode.
    ///
    /// # Errors
    ///
    /// Returns an error if no IRQ is available.
    pub fn allocate_irq(&self) -> Result<Irq> {
        let irq = self.next_irq.fetch_add(1, Ordering::SeqCst);
        if irq >= MAX_IRQS {
            return Err(crate::error::VmmError::Irq("IRQ exhausted".to_string()));
        }

        // Default: map IRQ N to GSI N for legacy compatibility
        let gsi = irq % MAX_GSIS;
        let config = IrqConfig {
            gsi,
            trigger_mode: TriggerMode::Edge,
            asserted: false,
        };

        {
            let mut configs = self
                .irq_configs
                .write()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            configs.insert(irq, config);
        }

        tracing::debug!("Allocated IRQ {} -> GSI {} (default edge)", irq, gsi);

        Ok(irq)
    }

    /// Configures an existing IRQ.
    pub fn configure_irq(&self, irq: Irq, gsi: Gsi, trigger_mode: TriggerMode) -> Result<()> {
        let mut configs = self
            .irq_configs
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(config) = configs.get_mut(&irq) {
            config.gsi = gsi;
            config.trigger_mode = trigger_mode;
            tracing::debug!(
                "Reconfigured IRQ {} -> GSI {}, mode={:?}",
                irq,
                gsi,
                trigger_mode
            );
            Ok(())
        } else {
            Err(crate::error::VmmError::Irq(format!(
                "IRQ {irq} not allocated"
            )))
        }
    }

    /// Triggers an interrupt.
    ///
    /// For edge-triggered IRQs, this sends a pulse (assert then deassert).
    /// For level-triggered IRQs, this asserts the line (use `deassert_irq` to clear).
    ///
    /// Implements interrupt coalescing: if the IRQ is already pending,
    /// the trigger is coalesced (not delivered again).
    ///
    /// # Errors
    ///
    /// Returns an error if the interrupt cannot be delivered.
    pub fn trigger_irq(&self, irq: Irq) -> Result<()> {
        // Check if masked
        if self.is_masked(irq) {
            tracing::trace!("IRQ {} is masked, not triggering", irq);
            return Ok(());
        }

        // Get configuration
        let configs = self
            .irq_configs
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let config = configs.get(&irq);
        let (gsi, trigger_mode) = match config {
            Some(c) => (c.gsi, c.trigger_mode),
            None => {
                // Legacy fallback: IRQ maps directly to GSI
                (irq % MAX_GSIS, TriggerMode::Edge)
            }
        };
        drop(configs);

        // Check for coalescing (only for edge-triggered)
        if trigger_mode == TriggerMode::Edge {
            let irq_bit = 1u32 << (irq % 32);
            let old_pending = self.pending.fetch_or(irq_bit, Ordering::SeqCst);
            if (old_pending & irq_bit) != 0 {
                // Already pending, coalesce
                self.stats.coalesced.fetch_add(1, Ordering::Relaxed);
                tracing::trace!("IRQ {} coalesced (already pending)", irq);
                return Ok(());
            }
        }

        self.stats.triggered.fetch_add(1, Ordering::Relaxed);
        tracing::trace!("Triggering IRQ {} -> GSI {}", irq, gsi);

        // Invoke the trigger callback
        let callback = self
            .trigger_callback
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(ref cb) = *callback {
            match trigger_mode {
                TriggerMode::Edge => {
                    // Edge: pulse the line (assert then deassert)
                    cb(gsi, true)?;
                    cb(gsi, false)?;
                    // Clear pending after delivery
                    let irq_bit = 1u32 << (irq % 32);
                    self.pending.fetch_and(!irq_bit, Ordering::SeqCst);
                }
                TriggerMode::Level => {
                    // Level: just assert (caller must deassert when done)
                    cb(gsi, true)?;
                    // Mark as asserted in config
                    drop(callback);
                    let mut configs = self
                        .irq_configs
                        .write()
                        .unwrap_or_else(std::sync::PoisonError::into_inner);
                    if let Some(c) = configs.get_mut(&irq) {
                        c.asserted = true;
                    }
                }
            }
        } else {
            tracing::warn!(
                "IRQ {} triggered but no callback registered (GSI {})",
                irq,
                gsi
            );
        }

        Ok(())
    }

    /// Deasserts a level-triggered interrupt.
    ///
    /// For level-triggered IRQs, the device calls this when the interrupt
    /// condition is cleared (e.g., data read from FIFO).
    pub fn deassert_irq(&self, irq: Irq) -> Result<()> {
        let configs = self
            .irq_configs
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let config = configs.get(&irq);
        let gsi = match config {
            Some(c) if c.trigger_mode == TriggerMode::Level => c.gsi,
            Some(_) => {
                tracing::trace!("deassert_irq called on edge-triggered IRQ {}", irq);
                return Ok(());
            }
            None => {
                return Ok(());
            }
        };
        drop(configs);

        // Invoke callback to deassert
        let callback = self
            .trigger_callback
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(ref cb) = *callback {
            cb(gsi, false)?;
        }

        // Mark as deasserted
        let mut configs = self
            .irq_configs
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(c) = configs.get_mut(&irq) {
            c.asserted = false;
        }

        tracing::trace!("Deasserted IRQ {} (GSI {})", irq, gsi);

        Ok(())
    }

    /// Acknowledges an interrupt (clears pending state).
    ///
    /// Called by the interrupt handler after processing to allow
    /// new interrupts of the same type.
    pub fn ack_irq(&self, irq: Irq) {
        if irq < 32 {
            let irq_bit = 1u32 << irq;
            self.pending.fetch_and(!irq_bit, Ordering::SeqCst);
            tracing::trace!("Acknowledged IRQ {}", irq);
        }
    }

    /// Masks an interrupt.
    pub fn mask_irq(&self, irq: Irq) {
        if irq < 32 {
            let old = self.mask.fetch_or(1 << irq, Ordering::SeqCst);
            tracing::trace!("Masked IRQ {}, old mask: {:#x}", irq, old);
        }
    }

    /// Unmasks an interrupt.
    pub fn unmask_irq(&self, irq: Irq) {
        if irq < 32 {
            let old = self.mask.fetch_and(!(1 << irq), Ordering::SeqCst);
            tracing::trace!("Unmasked IRQ {}, old mask: {:#x}", irq, old);
        }
    }

    /// Checks if an IRQ is masked.
    #[must_use]
    pub fn is_masked(&self, irq: Irq) -> bool {
        if irq < 32 {
            (self.mask.load(Ordering::SeqCst) & (1 << irq)) != 0
        } else {
            false
        }
    }

    /// Checks if an IRQ is pending.
    #[must_use]
    pub fn is_pending(&self, irq: Irq) -> bool {
        if irq < 32 {
            (self.pending.load(Ordering::SeqCst) & (1 << irq)) != 0
        } else {
            false
        }
    }

    /// Gets the GSI for an IRQ.
    #[must_use]
    pub fn get_gsi(&self, irq: Irq) -> Option<Gsi> {
        let configs = self
            .irq_configs
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        configs.get(&irq).map(|c| c.gsi)
    }

    /// Gets the trigger mode for an IRQ.
    #[must_use]
    pub fn get_trigger_mode(&self, irq: Irq) -> Option<TriggerMode> {
        let configs = self
            .irq_configs
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        configs.get(&irq).map(|c| c.trigger_mode)
    }

    /// Returns interrupt statistics.
    pub fn stats(&self) -> (u64, u64) {
        (
            self.stats.triggered.load(Ordering::Relaxed),
            self.stats.coalesced.load(Ordering::Relaxed),
        )
    }

    /// Resets statistics.
    pub fn reset_stats(&self) {
        self.stats.triggered.store(0, Ordering::Relaxed);
        self.stats.coalesced.store(0, Ordering::Relaxed);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;

    #[test]
    fn test_irq_allocation() {
        let chip = IrqChip::new().unwrap();

        let irq1 = chip.allocate_irq().unwrap();
        let irq2 = chip.allocate_irq().unwrap();

        assert!(irq2 > irq1);
    }

    #[test]
    fn test_irq_allocation_with_config() {
        let chip = IrqChip::new().unwrap();

        let irq = chip
            .allocate_irq_with_config(5, TriggerMode::Level)
            .unwrap();

        assert_eq!(chip.get_gsi(irq), Some(5));
        assert_eq!(chip.get_trigger_mode(irq), Some(TriggerMode::Level));
    }

    #[test]
    fn test_irq_masking() {
        let chip = IrqChip::new().unwrap();

        assert!(!chip.is_masked(0));

        chip.mask_irq(0);
        assert!(chip.is_masked(0));

        chip.unmask_irq(0);
        assert!(!chip.is_masked(0));
    }

    #[test]
    fn test_irq_device_trigger_chain() {
        let chip = Arc::new(IrqChip::new().unwrap());
        let events = Arc::new(Mutex::new(Vec::new()));
        let events_clone = Arc::clone(&events);

        let callback: IrqTriggerCallback = Box::new(move |gsi, level| {
            events_clone
                .lock()
                .unwrap_or_else(|e| e.into_inner())
                .push((gsi, level));
            Ok(())
        });
        chip.set_trigger_callback(Arc::new(callback));

        // Use a legacy IRQ to exercise pending tracking.
        let irq: Irq = 5;
        {
            let mut configs = chip.irq_configs.write().unwrap_or_else(|e| e.into_inner());
            configs.insert(
                irq,
                IrqConfig {
                    gsi: 5,
                    trigger_mode: TriggerMode::Edge,
                    asserted: false,
                },
            );
        }

        assert!(!chip.is_pending(irq));
        chip.trigger_irq(irq).unwrap();
        assert!(!chip.is_pending(irq));

        let recorded = events.lock().unwrap_or_else(|e| e.into_inner());
        assert_eq!(recorded.as_slice(), &[(5, true), (5, false)]);
    }

    #[test]
    fn test_irq_trigger_with_callback() {
        let chip = IrqChip::new().unwrap();
        let trigger_count = Arc::new(AtomicUsize::new(0));
        let count_clone = Arc::clone(&trigger_count);

        // Set up callback
        let callback: IrqTriggerCallback = Box::new(move |_gsi, _level| {
            count_clone.fetch_add(1, Ordering::SeqCst);
            Ok(())
        });
        chip.set_trigger_callback(Arc::new(callback));

        // Allocate and trigger
        let irq = chip.allocate_irq_with_config(1, TriggerMode::Edge).unwrap();
        chip.trigger_irq(irq).unwrap();

        // Edge-triggered should call callback twice (assert + deassert)
        assert_eq!(trigger_count.load(Ordering::SeqCst), 2);
    }

    #[test]
    fn test_irq_coalescing() {
        let chip = IrqChip::new().unwrap();
        let trigger_count = Arc::new(AtomicUsize::new(0));
        let count_clone = Arc::clone(&trigger_count);

        // Set up callback that doesn't clear pending
        let callback: IrqTriggerCallback = Box::new(move |_gsi, _level| {
            count_clone.fetch_add(1, Ordering::SeqCst);
            Ok(())
        });
        chip.set_trigger_callback(Arc::new(callback));

        // Allocate edge-triggered IRQ
        let irq = chip.allocate_irq_with_config(1, TriggerMode::Edge).unwrap();

        // Trigger multiple times
        chip.trigger_irq(irq).unwrap();
        chip.trigger_irq(irq).unwrap(); // Should be coalesced
        chip.trigger_irq(irq).unwrap(); // Should be coalesced

        // Only first should trigger (edge cleared pending immediately)
        // Actually for edge-triggered, pending is cleared after delivery
        // So subsequent triggers should also go through
        let (triggered, coalesced) = chip.stats();
        assert!(triggered >= 1);
        // The exact behavior depends on timing
        tracing::debug!("triggered={}, coalesced={}", triggered, coalesced);
    }

    #[test]
    fn test_level_triggered_irq() {
        let chip = IrqChip::new().unwrap();
        let levels = Arc::new(Mutex::new(Vec::new()));
        let levels_clone = Arc::clone(&levels);

        let callback: IrqTriggerCallback = Box::new(move |gsi, level| {
            levels_clone
                .lock()
                .unwrap_or_else(|e| e.into_inner())
                .push((gsi, level));
            Ok(())
        });
        chip.set_trigger_callback(Arc::new(callback));

        let irq = chip
            .allocate_irq_with_config(3, TriggerMode::Level)
            .unwrap();

        // Assert
        chip.trigger_irq(irq).unwrap();
        // Deassert
        chip.deassert_irq(irq).unwrap();

        let recorded = levels.lock().unwrap_or_else(|e| e.into_inner());
        assert_eq!(recorded.len(), 2);
        assert_eq!(recorded[0], (3, true)); // assert
        assert_eq!(recorded[1], (3, false)); // deassert
    }

    #[test]
    fn test_masked_irq_not_triggered() {
        let chip = IrqChip::new().unwrap();
        let trigger_count = Arc::new(AtomicUsize::new(0));
        let count_clone = Arc::clone(&trigger_count);

        let callback: IrqTriggerCallback = Box::new(move |_gsi, _level| {
            count_clone.fetch_add(1, Ordering::SeqCst);
            Ok(())
        });
        chip.set_trigger_callback(Arc::new(callback));

        // Allocate IRQ (will be >= 32, outside maskable range)
        let irq = chip.allocate_irq_with_config(0, TriggerMode::Edge).unwrap();

        // For this test, we need to test with a legacy IRQ that's maskable
        // Use mask_irq with IRQ 5 (a legacy IRQ)
        let legacy_irq: Irq = 5;

        // Insert config for legacy IRQ 5
        {
            let mut configs = chip.irq_configs.write().unwrap_or_else(|e| e.into_inner());
            configs.insert(
                legacy_irq,
                super::IrqConfig {
                    gsi: 5,
                    trigger_mode: TriggerMode::Edge,
                    asserted: false,
                },
            );
        }

        // Mask legacy IRQ and try to trigger
        chip.mask_irq(legacy_irq);
        chip.trigger_irq(legacy_irq).unwrap();

        // Should not have triggered because masked
        assert_eq!(trigger_count.load(Ordering::SeqCst), 0);

        // Unmask and trigger
        chip.unmask_irq(legacy_irq);
        chip.trigger_irq(legacy_irq).unwrap();

        // Should have triggered now
        assert_eq!(trigger_count.load(Ordering::SeqCst), 2); // assert + deassert

        // Also test that IRQs >= 32 are not affected by mask
        // IRQ was allocated with allocate_irq_with_config so it's >= 32
        assert!(irq >= 32);
        // These IRQs are never masked (is_masked returns false)
        assert!(!chip.is_masked(irq));
    }
}