vyre 0.4.0

// Program — a complete, self-contained GPU compute dispatch.
//
// A `Program` can be constructed without a GPU, serialized to disk,
// transmitted over a network, optimized by transformation passes,
// and lowered to any target backend. It is the unit of work in vyre.

use crate::ir::model::arena::{ArenaProgram, ExprArena};
use crate::ir::model::node::Node;
use crate::ir::model::types::{BufferAccess, DataType};
use rustc_hash::FxHashMap;
use std::collections::HashMap;
use std::hash::Hash;
use std::rc::Rc;
use std::sync::Arc;

/// A complete vyre program.
///
/// Contains everything needed to execute a GPU compute dispatch:
/// buffer declarations, workgroup configuration, and the entry point body.
///
/// # Example
///
/// A program that XORs two input buffers element-wise:
///
/// ```rust
/// use vyre::ir::{Program, BufferDecl, BufferAccess, DataType, Node, Expr, BinOp};
///
/// let program = Program::new(
///     vec![
///         BufferDecl::storage("a", 0, BufferAccess::ReadOnly, DataType::U32),
///         BufferDecl::storage("b", 1, BufferAccess::ReadOnly, DataType::U32),
///         BufferDecl::storage("out", 2, BufferAccess::ReadWrite, DataType::U32),
///     ],
///     [64, 1, 1],
///     vec![
///         Node::let_bind("idx", Expr::gid_x()),
///         Node::if_then(
///             Expr::lt(Expr::var("idx"), Expr::buf_len("out")),
///             vec![
///                 Node::store("out", Expr::var("idx"),
///                     Expr::bitxor(
///                         Expr::load("a", Expr::var("idx")),
///                         Expr::load("b", Expr::var("idx")),
///                     ),
///                 ),
///             ],
///         ),
///     ],
/// );
/// assert_eq!(program.buffers().len(), 3);
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Program {
    /// Stable ID of the certified operation this program implements.
    ///
    /// Runtime lowering must reject programs without an ID because anonymous IR
    /// cannot be tied back to a conform registry entry.
    pub entry_op_id: Option<String>,
    /// Buffer declarations. Each declares a named, typed, bound memory region.
    pub buffers: Arc<[BufferDecl]>,
    /// Sidecar index for O(1) buffer lookup by name.
    #[doc(hidden)]
    pub buffer_index: Arc<FxHashMap<String, usize>>,
    /// Workgroup size: `[x, y, z]`. Controls `@workgroup_size` in WGSL.
    pub workgroup_size: [u32; 3],
    /// Entry point body. Executes once per invocation.
    pub entry: Arc<Vec<Node>>,
}

/// A named buffer binding in a program.
///
/// # Examples
///
/// ```
/// use vyre::ir::{BufferDecl, BufferAccess, DataType};
///
/// let buf = BufferDecl::read("input", 0, DataType::U32);
/// assert_eq!(buf.name(), "input");
/// assert_eq!(buf.binding(), 0);
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct BufferDecl {
    /// Human-readable name. Referenced by `Expr::Load`, `Node::Store`, etc.
    pub name: String,
    /// Binding slot: `@binding(N)`. All buffers are in `@group(0)`.
    /// Ignored for `BufferAccess::Workgroup`.
    pub binding: u32,
    /// Access mode.
    pub access: BufferAccess,
    /// Element data type.
    pub element: DataType,
    /// Number of elements.
    ///
    /// For `Workgroup` memory this is the static array length.
    /// For storage and uniform buffers this is `0` (runtime-sized).
    pub count: u32,
    /// Whether this buffer is the scalar expression output for composition inlining.
    pub is_output: bool,
}

/// Refcount-shared, copy-on-write lexical scope map for IR validation.
///
/// `Scope` lets branch validators share an unchanged parent binding table and
/// allocate only when a branch actually inserts or updates a binding. Cloning a
/// scope is O(1); the first write after a clone copies the underlying map.
///
/// # Examples
///
/// ```
/// use vyre::ir::program::Scope;
///
/// let mut root = Scope::new();
/// root.insert("idx".to_string(), 0_u32);
///
/// let mut branch = root.child();
/// branch.insert("tmp".to_string(), 1_u32);
///
/// assert_eq!(root.get("idx"), Some(&0));
/// assert!(!root.contains_key("tmp"));
/// assert_eq!(branch.get("tmp"), Some(&1));
/// ```
#[derive(Debug, Clone)]
pub struct Scope<K, V> {
    bindings: Rc<HashMap<K, V>>,
}

impl<K, V> Default for Scope<K, V> {
    fn default() -> Self {
        Self {
            bindings: Rc::new(HashMap::new()),
        }
    }
}
mod impl_scope;

impl BufferDecl {
    /// Create a storage buffer declaration.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::{BufferDecl, BufferAccess, DataType};
    /// let _ = BufferDecl::storage("a", 0, BufferAccess::ReadOnly, DataType::U32);
    /// ```
    #[must_use]
    #[inline]
    pub fn storage(name: &str, binding: u32, access: BufferAccess, element: DataType) -> Self {
        Self {
            name: name.to_string(),
            binding,
            access,
            element,
            count: 0,
            is_output: false,
        }
    }

    /// Shorthand for a read-only storage buffer.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::{BufferDecl, DataType};
    /// let _ = BufferDecl::read("a", 0, DataType::U32);
    /// ```
    #[must_use]
    #[inline]
    pub fn read(name: &str, binding: u32, element: DataType) -> Self {
        Self::storage(name, binding, BufferAccess::ReadOnly, element)
    }

    /// Shorthand for a read-write storage buffer.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::{BufferDecl, DataType};
    /// let _ = BufferDecl::read_write("a", 0, DataType::U32);
    /// ```
    #[must_use]
    #[inline]
    pub fn read_write(name: &str, binding: u32, element: DataType) -> Self {
        Self::storage(name, binding, BufferAccess::ReadWrite, element)
    }

    /// Shorthand for the read-write result buffer used by call inlining.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::{BufferDecl, DataType};
    /// let _ = BufferDecl::output("a", 0, DataType::U32);
    /// ```
    #[must_use]
    #[inline]
    pub fn output(name: &str, binding: u32, element: DataType) -> Self {
        Self {
            is_output: true,
            ..Self::read_write(name, binding, element)
        }
    }

    /// Shorthand for a uniform buffer.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::{BufferDecl, DataType};
    /// let _ = BufferDecl::uniform("a", 0, DataType::U32);
    /// ```
    #[must_use]
    #[inline]
    pub fn uniform(name: &str, binding: u32, element: DataType) -> Self {
        Self::storage(name, binding, BufferAccess::Uniform, element)
    }

    /// Shorthand for a workgroup-local shared array.
    ///
    /// `count` is the static number of elements visible to all invocations
    /// in the same workgroup.
    #[must_use]
    #[inline]
    pub fn workgroup(name: &str, count: u32, element: DataType) -> Self {
        Self {
            name: name.to_string(),
            binding: 0,
            access: BufferAccess::Workgroup,
            element,
            count,
            is_output: false,
        }
    }

    /// Buffer name.
    #[must_use]
    #[inline]
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Binding slot.
    #[must_use]
    #[inline]
    pub fn binding(&self) -> u32 {
        self.binding
    }

    /// Buffer access mode.
    #[must_use]
    #[inline]
    pub fn access(&self) -> BufferAccess {
        self.access.clone()
    }

    /// Element data type.
    #[must_use]
    #[inline]
    pub fn element(&self) -> DataType {
        self.element.clone()
    }

    /// Static element count for workgroup buffers.
    #[must_use]
    #[inline]
    pub fn count(&self) -> u32 {
        self.count
    }

    /// Return true when this buffer is the unique inlining result buffer.
    #[must_use]
    #[inline]
    pub fn is_output(&self) -> bool {
        self.is_output
    }
}

impl Program {
    /// Create a complete program from buffer declarations, workgroup size, and
    /// entry-point nodes.
    #[must_use]
    #[inline]
    pub fn new(buffers: Vec<BufferDecl>, workgroup_size: [u32; 3], entry: Vec<Node>) -> Self {
        let buffer_index = Self::build_buffer_index(&buffers);
        Self {
            entry_op_id: None,
            buffers: Arc::from(buffers),
            buffer_index: Arc::new(buffer_index),
            workgroup_size,
            entry: Arc::new(entry),
        }
    }

    /// Create an arena-backed program scaffold.
    ///
    /// This constructor is the opt-in migration path for builders that want
    /// [`ExprRef`](crate::ir::model::arena::ExprRef) handles instead of boxed
    /// expression trees. [`Program::new`] remains the boxed-tree constructor.
    #[must_use]
    #[inline]
    pub fn with_arena<'a>(
        arena: &'a ExprArena,
        buffers: Vec<BufferDecl>,
        workgroup_size: [u32; 3],
    ) -> ArenaProgram<'a> {
        ArenaProgram::new(arena, buffers, workgroup_size)
    }

    /// Create a minimal program with no buffers and an empty body.
    #[must_use]
    #[inline]
    pub fn empty() -> Self {
        Self {
            entry_op_id: None,
            buffers: Arc::from([]),
            buffer_index: Arc::new(FxHashMap::default()),
            workgroup_size: [1, 1, 1],
            entry: Arc::new(Vec::new()),
        }
    }

    /// Attach the stable operation ID whose conform registry entry certifies
    /// this program for runtime lowering.
    #[must_use]
    #[inline]
    pub fn with_entry_op_id(mut self, op_id: impl Into<String>) -> Self {
        self.entry_op_id = Some(op_id.into());
        self
    }

    /// Stable operation ID required by the conform gate.
    #[must_use]
    #[inline]
    pub fn entry_op_id(&self) -> Option<&str> {
        self.entry_op_id.as_deref()
    }

    /// Attach an optional operation ID while preserving anonymous test IR.
    #[must_use]
    #[inline]
    pub(crate) fn with_optional_entry_op_id(mut self, op_id: Option<String>) -> Self {
        self.entry_op_id = op_id;
        self
    }

    /// Look up a buffer declaration by name.
    #[must_use]
    #[inline]
    pub fn buffer(&self, name: &str) -> Option<&BufferDecl> {
        self.buffer_index
            .get(name)
            .and_then(|&index| self.buffers.get(index))
    }

    /// Declared buffers.
    #[must_use]
    #[inline]
    pub fn buffers(&self) -> &[BufferDecl] {
        self.buffers.as_ref()
    }

    /// Workgroup dimensions.
    #[must_use]
    #[inline]
    pub fn workgroup_size(&self) -> [u32; 3] {
        self.workgroup_size
    }

    /// Set the workgroup dimensions in place. Used by harnesses that
    /// need to clone-and-rewrite a program's workgroup size for fallback
    /// dispatch — the alternative was to reconstruct the entire Program,
    /// which is unnecessarily expensive when only one field changes.
    #[inline]
    pub fn set_workgroup_size(&mut self, workgroup_size: [u32; 3]) {
        self.workgroup_size = workgroup_size;
    }

    /// Entry-point nodes.
    #[must_use]
    #[inline]
    pub fn entry(&self) -> &[Node] {
        self.entry.as_ref().as_slice()
    }

    /// Mutable entry-point nodes for transformation passes.
    #[must_use]
    #[inline]
    pub fn entry_mut(&mut self) -> &mut Vec<Node> {
        Arc::make_mut(&mut self.entry)
    }

    /// Check whether a named buffer exists.
    #[must_use]
    #[inline]
    pub fn has_buffer(&self, name: &str) -> bool {
        self.buffer_index.contains_key(name)
    }

    /// Number of declared buffers.
    #[must_use]
    #[inline]
    pub fn buffer_count(&self) -> usize {
        self.buffers.len()
    }

    #[inline]
    fn build_buffer_index(buffers: &[BufferDecl]) -> FxHashMap<String, usize> {
        let mut index = FxHashMap::default();
        index.reserve(buffers.len());
        for (buffer_index, buffer) in buffers.iter().enumerate() {
            index.entry(buffer.name.clone()).or_insert(buffer_index);
        }
        index
    }
}