vyre 0.4.0

use crate::ir::DataType;
use crate::ops::{AlgebraicLaw, Backend, IntrinsicDescriptor, OpSpec};

pub const INPUTS: &[DataType] = &[DataType::U32, DataType::U32];

pub const OUTPUTS: &[DataType] = &[DataType::U32, DataType::U32];

pub const LAWS: &[AlgebraicLaw] = &[];

/// Registered op spec for `workgroup.visitor`.
pub const SPEC: OpSpec = OpSpec::intrinsic(
    "workgroup.visitor",
    INPUTS,
    OUTPUTS,
    LAWS,
    wgsl_only,
    IntrinsicDescriptor::new(
        "workgroup_visitor_visit",
        "workgroup-sram-mark-and-emit",
        crate::ops::cpu_op::structured_intrinsic_cpu,
    ),
);

/// Status word returned by `visit`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VisitStatus {
    /// First time this node has been visited. The caller may emit
    /// downstream effects.
    FirstVisit = 0,
    /// Node was already visited; the caller should skip emission.
    AlreadyVisited = 1,
}

/// Error returned by fallible visitor operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VisitError {
    /// The node id was beyond the declared universe.
    OutOfRange,
}

/// Bounded mark-and-emit visitor used as the CPU reference for
/// `workgroup.visitor`.
///
/// The visitor maintains a fixed-size bitmap of "visited" marks and
/// an emit buffer that records first-visits in encounter order. The
/// GPU lowering uses an atomic word-compare-and-swap over the same
/// bitmap so concurrent lanes cannot double-emit a node.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WorkgroupVisitor {
    universe: usize,
    emit_capacity: usize,
    visited: Vec<bool>,
    emit_buffer: Vec<u32>,
}

impl WorkgroupVisitor {
    /// Create an empty visitor over a universe of `universe` nodes
    /// with an emit buffer of `emit_capacity` slots.
    #[must_use]
    pub fn new(universe: usize, emit_capacity: usize) -> Self {
        Self {
            universe,
            emit_capacity,
            visited: vec![false; universe],
            emit_buffer: Vec::with_capacity(emit_capacity),
        }
    }

    /// Size of the node universe.
    #[must_use]
    pub fn universe(&self) -> usize {
        self.universe
    }

    /// Maximum emit-buffer length.
    #[must_use]
    pub fn emit_capacity(&self) -> usize {
        self.emit_capacity
    }

    /// Number of distinct first-visits recorded so far.
    #[must_use]
    pub fn emitted(&self) -> &[u32] {
        &self.emit_buffer
    }

    /// Whether the visitor has seen `node`.
    #[must_use]
    pub fn is_visited(&self, node: u32) -> bool {
        let idx = node as usize;
        idx < self.visited.len() && self.visited[idx]
    }

    /// Mark `node` as visited. Returns [`VisitStatus::FirstVisit`] if
    /// `node` was not previously marked; otherwise
    /// [`VisitStatus::AlreadyVisited`].
    ///
    /// # Errors
    ///
    /// Returns [`VisitError::OutOfRange`] when `node >= universe`.
    pub fn visit(&mut self, node: u32) -> Result<VisitStatus, VisitError> {
        let idx = node as usize;
        if idx >= self.universe {
            return Err(VisitError::OutOfRange);
        }
        if self.visited[idx] {
            return Ok(VisitStatus::AlreadyVisited);
        }
        self.visited[idx] = true;
        // Emit buffer silently truncates at capacity — the caller
        // can detect a missed emit by observing `emitted().len()`
        // plateauing across successive first-visits. This matches
        // the WGSL lowering's bounded shared-memory contract.
        if self.emit_buffer.len() < self.emit_capacity {
            self.emit_buffer.push(node);
        }
        Ok(VisitStatus::FirstVisit)
    }

    /// Clear all marks and the emit buffer. Cheap O(n) reset — the
    /// backing allocations are preserved for the next cycle.
    pub fn reset(&mut self) {
        for slot in &mut self.visited {
            *slot = false;
        }
        self.emit_buffer.clear();
    }
}

pub fn wgsl_only(backend: &Backend) -> bool {
    matches!(backend, Backend::Wgsl)
}