wgsl-rs 0.1.0-beta.4

//! Extension trait for downstream crates to modify WGSL IR after transpilation
//! but before rendering/instantiation.
//!
//! Extensions are specified via the `extensions` attribute on the `#[wgsl]`
//! macro:
//!
//! ```ignore
//! #[wgsl(crate_path = crate, extensions = [my_crate::MyExt])]
//! pub mod my_shader { ... }
//! ```
//!
//! Each extension type must implement [`WgslExtension`]. The
//! [`modify_ir`](WgslExtension::modify_ir) method is called on the
//! base/template IR module **before** type instantiation, so modifications made
//! by extensions are carried through every `instantiate()` call.
//!
//! # Statement macro lowering
//!
//! When the `#[wgsl]` parser encounters a statement macro that is not one of
//! its builtins (`slab_copy!`, `discard!`), it
//! passes it through as a [`crate::ir::Stmt::Macro`] variant instead of
//! rejecting it. An extension can then recognize the macro by name in
//! `modify_ir` and replace the `Stmt::Macro` with lowered IR.
//!
//! Extensions declare which macro names they handle via the
//! [`MACROS`](WgslExtension::MACROS) associated const. The `#[wgsl]` macro
//! emits a compile-time `const` check ensuring every `Stmt::Macro` in the
//! module is claimed by at least one listed extension — so a typo or
//! missing extension is a compile error, not a runtime error.

/// A trait for modifying WGSL IR modules after transpilation.
///
/// Extensions can inspect IR items (including their preserved Rust attributes)
/// and add, remove, or modify items. The IR module is passed mutably, so
/// extensions can make any changes they need.
///
/// # Statement macro lowering
///
/// Extensions that lower statement macros (see [`crate::ir::Stmt::Macro`])
/// should declare the macro names they handle in [`MACROS`](Self::MACROS). The
/// `#[wgsl]` macro reads this const at compile time to verify every
/// `Stmt::Macro` in the module is claimed by at least one listed extension.
///
/// # Example
///
/// ```ignore
/// use wgsl_rs::ir;
/// use wgsl_rs::WgslExtension;
///
/// pub struct SlabItemExt;
///
/// impl WgslExtension for SlabItemExt {
///     const MACROS: &'static [&'static str] = &["slab_read", "slab_write"];
///
///     fn modify_ir(module: &mut ir::Module) {
///         // Find structs with #[derive(SlabItem)] in their attrs
///         let slab_structs: Vec<(String, Vec<ir::Field>)> = module.items.iter()
///             .filter_map(|item| match item {
///                 ir::Item::Struct(s) if has_derive(&s.attrs, "SlabItem") => {
///                     Some((s.name.clone(), s.fields.clone()))
///                 }
///                 _ => None,
///             })
///             .collect();
///
///         // Add helper functions for each slab struct
///         for (name, fields) in slab_structs {
///             module.items.push(ir::Item::Fn(make_slab_read_fn(&name, &fields)));
///             module.items.push(ir::Item::Fn(make_slab_write_fn(&name, &fields)));
///         }
///     }
/// }
/// ```
pub trait WgslExtension {
    /// Statement-macro names this extension claims to handle in
    /// [`modify_ir`](Self::modify_ir).
    ///
    /// The `#[wgsl]` macro emits a compile-time `const` check ensuring every
    /// [`crate::ir::Stmt::Macro`] in the module is claimed by at least one
    /// listed extension. If a macro name is used in the module but no
    /// listed extension declares it in `MACROS`, the result is a compile
    /// error (`E0080`), not a runtime error.
    ///
    /// Extensions that do not lower statement macros can leave this as the
    /// empty default.
    const MACROS: &'static [&'static str] = &[];

    /// Modify the IR module in place.
    ///
    /// This method is called in the IR constructor generated by the `#[wgsl]`
    /// macro, after the module's IR items have been constructed but before
    /// type instantiation (for template modules).
    fn modify_ir(module: &mut crate::ir::Module);
}

// ---------------------------------------------------------------------------
// Compile-time macro-claim verification helpers
// ---------------------------------------------------------------------------

/// Check if a `&str` slice contains a given needle, byte-by-byte.
///
/// This is a `const fn` because `str::PartialEq` is not yet stable in
/// `const` context. The `#[wgsl]` macro emits `const _: () = { ... }`
/// blocks that call this to verify every `Stmt::Macro` name is claimed by
/// at least one listed extension's `MACROS` const.
pub const fn contains_str(haystack: &[&str], needle: &str) -> bool {
    let nb = needle.as_bytes();
    let mut i = 0;
    while i < haystack.len() {
        let hb = haystack[i].as_bytes();
        if hb.len() == nb.len() {
            let mut j = 0;
            let mut eq = true;
            while j < nb.len() {
                if hb[j] != nb[j] {
                    eq = false;
                }
                j += 1;
            }
            if eq {
                return true;
            }
        }
        i += 1;
    }
    false
}

/// Check if any extension's `MACROS` list contains the given macro name.
///
/// `exts` is a slice of `MACROS` slices — one per extension listed in
/// `extensions = [...]`. Returns `true` if at least one extension claims
/// the macro.
pub const fn any_ext_claims(exts: &[&[&str]], name: &str) -> bool {
    let mut i = 0;
    while i < exts.len() {
        if contains_str(exts[i], name) {
            return true;
        }
        i += 1;
    }
    false
}