vyre 0.4.0

// Expression nodes — produce values.
//
// Every expression evaluates to a typed value. Expressions are pure:
// they read state but do not modify it.

use crate::ir::model::types::{AtomicOp, BinOp, DataType, UnOp};
use std::borrow::Borrow;
use std::fmt;
use std::ops::Deref;
use std::sync::Arc;

/// Interned identifier used by expression nodes.
///
/// `Ident` is cheap to clone and keeps expression trees from repeatedly
/// allocating owned `String` values for the same variable or buffer names.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Ident(Arc<str>);

impl Ident {
    /// Return the identifier text.
    #[must_use]
    #[inline]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl From<&str> for Ident {
    #[inline]
    fn from(value: &str) -> Self {
        Self(Arc::from(value))
    }
}

impl From<String> for Ident {
    #[inline]
    fn from(value: String) -> Self {
        Self(Arc::from(value))
    }
}

impl Deref for Ident {
    type Target = str;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_str()
    }
}

impl AsRef<str> for Ident {
    #[inline]
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl Borrow<str> for Ident {
    #[inline]
    fn borrow(&self) -> &str {
        self.as_str()
    }
}

impl fmt::Display for Ident {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl PartialEq<str> for Ident {
    #[inline]
    fn eq(&self, other: &str) -> bool {
        self.as_str() == other
    }
}

impl PartialEq<&str> for Ident {
    #[inline]
    fn eq(&self, other: &&str) -> bool {
        self.as_str() == *other
    }
}

/// An expression that produces a value.
///
/// # Examples
///
/// ```
/// use vyre::ir::Expr;
///
/// let lit = Expr::u32(42);
/// let var = Expr::var("x");
/// let add = Expr::add(lit, var);
/// ```
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Literal unsigned 32-bit integer.
    LitU32(u32),

    /// Literal signed 32-bit integer.
    LitI32(i32),

    /// Literal IEEE 754 binary32 floating-point.
    LitF32(f32),

    /// Literal boolean.
    LitBool(bool),

    /// Local variable reference by name.
    Var(Ident),

    /// Read one element from a named buffer at the given index.
    Load {
        /// Buffer name (must match a `BufferDecl::name` in the `Program`).
        buffer: Ident,
        /// Index expression (element offset, not byte offset).
        index: Box<Expr>,
    },

    /// Element count of a named buffer (`arrayLength`).
    BufLen {
        /// Buffer name.
        buffer: Ident,
    },

    /// Global invocation ID component.
    InvocationId {
        /// Axis: 0 = x, 1 = y, 2 = z.
        axis: u8,
    },

    /// Workgroup ID component.
    WorkgroupId {
        /// Axis: 0 = x, 1 = y, 2 = z.
        axis: u8,
    },

    /// Local invocation ID within the workgroup.
    LocalId {
        /// Axis: 0 = x, 1 = y, 2 = z.
        axis: u8,
    },

    /// Binary operation.
    BinOp {
        /// Operator.
        op: BinOp,
        /// Left operand.
        left: Box<Expr>,
        /// Right operand.
        right: Box<Expr>,
    },

    /// Unary operation.
    UnOp {
        /// Operator.
        op: UnOp,
        /// Operand.
        operand: Box<Expr>,
    },

    /// Call another operation by ID.
    Call {
        /// Operation identifier to invoke.
        op_id: String,
        /// Arguments.
        args: Vec<Expr>,
    },

    /// Ternary select: `select(false_val, true_val, cond)`.
    Select {
        /// Condition.
        cond: Box<Expr>,
        /// Value when true.
        true_val: Box<Expr>,
        /// Value when false.
        false_val: Box<Expr>,
    },

    /// Type cast.
    Cast {
        /// Target type.
        target: DataType,
        /// Value to cast.
        value: Box<Expr>,
    },

    /// Fused multiply-add: `a * b + c` (f32).
    Fma {
        /// First operand.
        a: Box<Expr>,
        /// Second operand.
        b: Box<Expr>,
        /// Third operand.
        c: Box<Expr>,
    },

    /// Atomic buffer operation. Returns the value before the operation.
    Atomic {
        /// Atomic operation kind.
        op: AtomicOp,
        /// Buffer name (must be `ReadWrite`).
        buffer: Ident,
        /// Element index.
        index: Box<Expr>,
        /// Expected old value for compare-exchange. Must be `Some` only for
        /// `AtomicOp::CompareExchange`.
        expected: Option<Box<Expr>>,
        /// Value operand. For compare-exchange this is the replacement value.
        value: Box<Expr>,
    },
}

impl Expr {
    /// Load from buffer at index.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::Expr;
    /// let _ = Expr::load("a", Expr::u32(0));
    /// ```
    #[must_use]
    #[inline]
    pub fn load(buffer: &str, index: Self) -> Self {
        Self::Load {
            buffer: Ident::from(buffer),
            index: Box::new(index),
        }
    }

    /// Buffer element count.
    ///
    /// # Examples
    ///
    /// ```
    /// use vyre::ir::Expr;
    /// let _ = Expr::buf_len("a");
    /// ```
    #[must_use]
    #[inline]
    pub fn buf_len(buffer: &str) -> Self {
        Self::BufLen {
            buffer: Ident::from(buffer),
        }
    }

    /// `global_invocation_id.x`
    #[must_use]
    #[inline]
    pub fn gid_x() -> Self {
        Self::InvocationId { axis: 0 }
    }

    /// `global_invocation_id.y`
    #[must_use]
    #[inline]
    pub fn gid_y() -> Self {
        Self::InvocationId { axis: 1 }
    }

    /// `global_invocation_id.z`
    #[must_use]
    #[inline]
    pub fn gid_z() -> Self {
        Self::InvocationId { axis: 2 }
    }

    /// `workgroup_id.x`
    #[must_use]
    #[inline]
    pub fn workgroup_x() -> Self {
        Self::WorkgroupId { axis: 0 }
    }

    /// `workgroup_id.y`
    #[must_use]
    #[inline]
    pub fn workgroup_y() -> Self {
        Self::WorkgroupId { axis: 1 }
    }

    /// `workgroup_id.z`
    #[must_use]
    #[inline]
    pub fn workgroup_z() -> Self {
        Self::WorkgroupId { axis: 2 }
    }

    /// `local_invocation_id.x`
    #[must_use]
    #[inline]
    pub fn local_x() -> Self {
        Self::LocalId { axis: 0 }
    }

    /// `local_invocation_id.y`
    #[must_use]
    #[inline]
    pub fn local_y() -> Self {
        Self::LocalId { axis: 1 }
    }

    /// `local_invocation_id.z`
    #[must_use]
    #[inline]
    pub fn local_z() -> Self {
        Self::LocalId { axis: 2 }
    }

    /// Conditional select.
    #[must_use]
    #[inline]
    pub fn select(cond: Self, true_val: Self, false_val: Self) -> Self {
        Self::Select {
            cond: Box::new(cond),
            true_val: Box::new(true_val),
            false_val: Box::new(false_val),
        }
    }

    /// Named variable reference.
    #[must_use]
    #[inline]
    pub fn var(name: &str) -> Self {
        Self::Var(Ident::from(name))
    }

    /// Unsigned 32-bit literal.
    #[must_use]
    #[inline]
    pub fn u32(value: u32) -> Self {
        Self::LitU32(value)
    }

    /// Signed 32-bit literal.
    #[must_use]
    #[inline]
    pub fn i32(value: i32) -> Self {
        Self::LitI32(value)
    }

    /// 32-bit floating-point literal.
    #[must_use]
    #[inline]
    pub fn f32(value: f32) -> Self {
        Self::LitF32(value)
    }

    /// Boolean literal.
    #[must_use]
    #[inline]
    pub fn bool(value: bool) -> Self {
        Self::LitBool(value)
    }

    /// Operation call by stable operation ID.
    #[must_use]
    #[inline]
    pub fn call(op_id: &str, args: Vec<Self>) -> Self {
        Self::Call {
            op_id: op_id.to_string(),
            args,
        }
    }

    /// Fused multiply-add `a * b + c` (f32).
    #[must_use]
    #[inline]
    pub fn fma(a: Self, b: Self, c: Self) -> Self {
        Self::Fma {
            a: Box::new(a),
            b: Box::new(b),
            c: Box::new(c),
        }
    }

    /// Cast a value to `target`.
    #[must_use]
    #[inline]
    pub fn cast(target: DataType, value: Self) -> Self {
        Self::Cast {
            target,
            value: Box::new(value),
        }
    }
}
mod atomics;
mod builders;

#[cfg(test)]
mod tests {
    use super::Expr;

    #[test]
    fn expr_size_is_bounded() {
        let size = std::mem::size_of::<Expr>();
        eprintln!("Expr size: {size}");
        assert!(
            size <= 128,
            "Expr grew to {size} bytes. Fix: box the largest variant before adding more fields."
        );
    }
}