From 987430762cfe9668024946301f0e757c3e0b0279 Mon Sep 17 00:00:00 2001 From: firestar99 Date: Sun, 31 Aug 2025 12:30:07 +0200 Subject: [PATCH] shaders: add trait `BufferStruct` and derive macro --- Cargo.lock | 4 + Cargo.toml | 6 +- node-graph/gcore-shaders/Cargo.toml | 7 +- node-graph/gcore-shaders/src/blending.rs | 6 +- .../gcore-shaders/src/color/color_types.rs | 3 +- .../src/shaders/buffer_struct/glam.rs | 114 ++++++++ .../src/shaders/buffer_struct/mod.rs | 63 +++++ .../src/shaders/buffer_struct/primitive.rs | 135 +++++++++ node-graph/gcore-shaders/src/shaders/mod.rs | 5 + node-graph/gcore/src/lib.rs | 1 + node-graph/graster-nodes/Cargo.toml | 1 + node-graph/graster-nodes/src/adjustments.rs | 9 +- node-graph/node-macro/src/buffer_struct.rs | 261 ++++++++++++++++++ node-graph/node-macro/src/lib.rs | 9 + .../src/shader_nodes/per_pixel_adjust.rs | 21 +- .../per_pixel_adjust_runtime.rs | 6 +- 16 files changed, 631 insertions(+), 20 deletions(-) create mode 100644 node-graph/gcore-shaders/src/shaders/buffer_struct/glam.rs create mode 100644 node-graph/gcore-shaders/src/shaders/buffer_struct/mod.rs create mode 100644 node-graph/gcore-shaders/src/shaders/buffer_struct/primitive.rs create mode 100644 node-graph/node-macro/src/buffer_struct.rs diff --git a/Cargo.lock b/Cargo.lock index 4efc07dfe0..bad343ee50 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2136,8 +2136,10 @@ dependencies = [ "graphene-core", "half", "log", + "node-macro", "num-derive", "num-traits", + "num_enum", "serde", "specta", "spirv-std", @@ -2186,6 +2188,7 @@ dependencies = [ "ndarray", "node-macro", "num-traits", + "num_enum", "rand 0.9.2", "rand_chacha 0.9.0", "serde", @@ -5524,6 +5527,7 @@ version = "0.9.0" source = "git+https://github.com/rust-gpu/rust-gpu?rev=c12f216121820580731440ee79ebc7403d6ea04f#c12f216121820580731440ee79ebc7403d6ea04f" dependencies = [ "bitflags 1.3.2", + "bytemuck", "glam", "libm", "num-traits", diff --git a/Cargo.toml b/Cargo.toml index ca1a6ed5df..b22749be3b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,7 +84,7 @@ graphite-proc-macros = { path = "proc-macros" } # Workspace dependencies rustc-hash = "2.0" -bytemuck = { version = "1.13", features = ["derive"] } +bytemuck = { version = "1.13", features = ["derive", "min_const_generics"] } serde = { version = "1.0", features = ["derive", "rc"] } serde_json = "1.0" serde-wasm-bindgen = "0.6" @@ -154,7 +154,7 @@ parley = "0.5" skrifa = "0.36" pretty_assertions = "1.4" fern = { version = "0.7", features = ["colored"] } -num_enum = "0.7" +num_enum = { version = "0.7", default-features = false } num-derive = "0.4" num-traits = { version = "0.2", default-features = false, features = ["libm"] } specta = { version = "2.0.0-rc.22", features = [ @@ -193,7 +193,7 @@ open = "5.3" poly-cool = "0.3" spin = "0.10" clap = "4.5" -spirv-std = { git = "https://github.com/rust-gpu/rust-gpu", rev = "c12f216121820580731440ee79ebc7403d6ea04f" } +spirv-std = { git = "https://github.com/rust-gpu/rust-gpu", rev = "c12f216121820580731440ee79ebc7403d6ea04f", features = ["bytemuck"] } cargo-gpu = { git = "https://github.com/rust-gpu/cargo-gpu", rev = "f969528e87baa17a7d48eecf4a6fcfdcaaf30566" } [workspace.lints.rust] diff --git a/node-graph/gcore-shaders/Cargo.toml b/node-graph/gcore-shaders/Cargo.toml index 6e38c66386..b9a37d24c8 100644 --- a/node-graph/gcore-shaders/Cargo.toml +++ b/node-graph/gcore-shaders/Cargo.toml @@ -22,10 +22,14 @@ std = [ "glam/serde", "half/std", "half/serde", - "num-traits/std" + "num-traits/std", + "num_enum/std", ] [dependencies] +# Local dependencies +node-macro = { workspace = true } + # Local std dependencies dyn-any = { workspace = true, optional = true } @@ -35,6 +39,7 @@ glam = { workspace = true } half = { workspace = true, default-features = false } num-derive = { workspace = true } num-traits = { workspace = true } +num_enum = { workspace = true } spirv-std = { workspace = true } # Workspace std dependencies diff --git a/node-graph/gcore-shaders/src/blending.rs b/node-graph/gcore-shaders/src/blending.rs index b305dd0910..747c4ba0a8 100644 --- a/node-graph/gcore-shaders/src/blending.rs +++ b/node-graph/gcore-shaders/src/blending.rs @@ -1,9 +1,11 @@ use core::fmt::Display; use core::hash::{Hash, Hasher}; +use node_macro::BufferStruct; +use num_enum::{FromPrimitive, IntoPrimitive}; #[cfg(not(feature = "std"))] use num_traits::float::Float; -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, Copy, PartialEq, BufferStruct)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "std", serde(default))] pub struct AlphaBlending { @@ -66,7 +68,7 @@ impl AlphaBlending { } #[repr(i32)] -#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, bytemuck::NoUninit)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, BufferStruct, FromPrimitive, IntoPrimitive)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] pub enum BlendMode { // Basic group diff --git a/node-graph/gcore-shaders/src/color/color_types.rs b/node-graph/gcore-shaders/src/color/color_types.rs index 63517341e8..5127a0c07a 100644 --- a/node-graph/gcore-shaders/src/color/color_types.rs +++ b/node-graph/gcore-shaders/src/color/color_types.rs @@ -5,6 +5,7 @@ use core::fmt::Debug; use core::hash::Hash; use glam::Vec4; use half::f16; +use node_macro::BufferStruct; #[cfg(not(feature = "std"))] use num_traits::Euclid; #[cfg(not(feature = "std"))] @@ -215,7 +216,7 @@ impl Pixel for Luma {} /// The other components (RGB) are stored as `f32` that range from `0.0` up to `f32::MAX`, /// the values encode the brightness of each channel proportional to the light intensity in cd/m² (nits) in HDR, and `0.0` (black) to `1.0` (white) in SDR color. #[repr(C)] -#[derive(Debug, Default, Clone, Copy, PartialEq, Pod, Zeroable)] +#[derive(Debug, Default, Clone, Copy, PartialEq, Pod, Zeroable, BufferStruct)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] pub struct Color { red: f32, diff --git a/node-graph/gcore-shaders/src/shaders/buffer_struct/glam.rs b/node-graph/gcore-shaders/src/shaders/buffer_struct/glam.rs new file mode 100644 index 0000000000..568f017045 --- /dev/null +++ b/node-graph/gcore-shaders/src/shaders/buffer_struct/glam.rs @@ -0,0 +1,114 @@ +use crate::shaders::buffer_struct::BufferStruct; + +macro_rules! glam_array { + ($t:ty, $a:ty) => { + unsafe impl BufferStruct for $t { + type Buffer = $a; + + #[inline] + fn write(from: Self) -> Self::Buffer { + <$t>::to_array(&from) + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + <$t>::from_array(from) + } + } + }; +} + +macro_rules! glam_cols_array { + ($t:ty, $a:ty) => { + unsafe impl BufferStruct for $t { + type Buffer = $a; + + #[inline] + fn write(from: Self) -> Self::Buffer { + <$t>::to_cols_array(&from) + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + <$t>::from_cols_array(&from) + } + } + }; +} + +glam_array!(glam::Vec2, [f32; 2]); +glam_array!(glam::Vec3, [f32; 3]); +// glam_array!(Vec3A, [f32; 4]); +glam_array!(glam::Vec4, [f32; 4]); +glam_array!(glam::Quat, [f32; 4]); +glam_cols_array!(glam::Mat2, [f32; 4]); +glam_cols_array!(glam::Mat3, [f32; 9]); +// glam_cols_array!(Mat3A, [f32; 4]); +glam_cols_array!(glam::Mat4, [f32; 16]); +glam_cols_array!(glam::Affine2, [f32; 6]); +glam_cols_array!(glam::Affine3A, [f32; 12]); + +glam_array!(glam::DVec2, [f64; 2]); +glam_array!(glam::DVec3, [f64; 3]); +glam_array!(glam::DVec4, [f64; 4]); +glam_array!(glam::DQuat, [f64; 4]); +glam_cols_array!(glam::DMat2, [f64; 4]); +glam_cols_array!(glam::DMat3, [f64; 9]); +glam_cols_array!(glam::DMat4, [f64; 16]); +glam_cols_array!(glam::DAffine2, [f64; 6]); +glam_cols_array!(glam::DAffine3, [f64; 12]); + +glam_array!(glam::I16Vec2, [i16; 2]); +glam_array!(glam::I16Vec3, [i16; 3]); +glam_array!(glam::I16Vec4, [i16; 4]); + +glam_array!(glam::U16Vec2, [u16; 2]); +glam_array!(glam::U16Vec3, [u16; 3]); +glam_array!(glam::U16Vec4, [u16; 4]); + +glam_array!(glam::IVec2, [i32; 2]); +glam_array!(glam::IVec3, [i32; 3]); +glam_array!(glam::IVec4, [i32; 4]); + +glam_array!(glam::UVec2, [u32; 2]); +glam_array!(glam::UVec3, [u32; 3]); +glam_array!(glam::UVec4, [u32; 4]); + +glam_array!(glam::I64Vec2, [i64; 2]); +glam_array!(glam::I64Vec3, [i64; 3]); +glam_array!(glam::I64Vec4, [i64; 4]); + +glam_array!(glam::U64Vec2, [u64; 2]); +glam_array!(glam::U64Vec3, [u64; 3]); +glam_array!(glam::U64Vec4, [u64; 4]); + +unsafe impl BufferStruct for glam::Vec3A { + type Buffer = [f32; 4]; + + #[inline] + fn write(from: Self) -> Self::Buffer { + glam::Vec4::to_array(&from.extend(0.)) + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + glam::Vec3A::from_vec4(glam::Vec4::from_array(from)) + } +} + +/// do NOT use slices, otherwise spirv will fail to compile +unsafe impl BufferStruct for glam::Mat3A { + type Buffer = [f32; 12]; + + #[inline] + fn write(from: Self) -> Self::Buffer { + let a = from.to_cols_array(); + [a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], 0., 0., 0.] + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + let a = from; + glam::Mat3A::from_cols_array(&[a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8]]) + } +} diff --git a/node-graph/gcore-shaders/src/shaders/buffer_struct/mod.rs b/node-graph/gcore-shaders/src/shaders/buffer_struct/mod.rs new file mode 100644 index 0000000000..07fca85e46 --- /dev/null +++ b/node-graph/gcore-shaders/src/shaders/buffer_struct/mod.rs @@ -0,0 +1,63 @@ +//! I (@firestar99) copied this entire mod from one of my projects, as I haven't uploaded that lib to crates. Hopefully +//! rust-gpu improves and this entire thing becomes unnecessary in the future. +//! +//! https://github.com/Firestar99/nanite-at-home/tree/008dac8df656959c71efeddd2d3ddabcb801771c/rust-gpu-bindless/crates/buffer-content + +use bytemuck::Pod; + +mod glam; +mod primitive; + +/// A BufferStruct is a "parallel representation" of the original struct with some fundamental types remapped. This +/// struct hierarchy represents how data is stored in GPU Buffers, where all types must be [`Pod`] to allow +/// transmuting them to `&[u8]` with [`bytemuck`]. +/// +/// Notable type remappings (original: buffer): +/// * bool: u32 of 0 or 1 +/// * any repr(u32) enum: u32 with remapping via [`num_enum`] +/// +/// By adding `#[derive(ShaderStruct)]` to your struct (or enum), a parallel `{name}Buffer` struct is created with all +/// the members of the original struct, but with their types using the associated remapped types as specified by this +/// trait. +/// +/// # Origin +/// I (@firestar99) copied this entire mod from my [Nanite-at-home] project, specifically the [buffer-content] crate +/// and the [buffer_struct] proc macro. The variant here has quite some modifications, to both cleaned up some of the +/// mistakes my implementation has and to customize it a bit for graphite. +/// +/// Hopefully rust-gpu improves to the point where this remapping becomes unnecessary. +/// +/// [Nanite-at-home]: https://github.com/Firestar99/nanite-at-home +/// [buffer-content]: https://github.com/Firestar99/nanite-at-home/tree/008dac8df656959c71efeddd2d3ddabcb801771c/rust-gpu-bindless/crates/buffer-content +/// [buffer_struct]: https://github.com/Firestar99/nanite-at-home/blob/008dac8df656959c71efeddd2d3ddabcb801771c/rust-gpu-bindless/crates/macros/src/buffer_struct.rs +/// +/// # Safety +/// The associated type Transfer must be the same on all targets. Writing followed by reading back a value must result +/// in the same value. +pub unsafe trait BufferStruct: Copy + Send + Sync + 'static { + type Buffer: Pod + Send + Sync; + + fn write(from: Self) -> Self::Buffer; + + fn read(from: Self::Buffer) -> Self; +} + +/// Trait marking all [`BufferStruct`] whose read and write methods are identity. While [`BufferStruct`] only +/// requires `t == read(write(t))`, this trait additionally requires `t == read(t) == write(t)`. As this removes the +/// conversion requirement for writing to or reading from a buffer, one can acquire slices from buffers created of these +/// types. +/// +/// Implementing this type is completely safe due to the [`Pod`] requirement. +pub trait BufferStructIdentity: Pod + Send + Sync {} + +unsafe impl BufferStruct for T { + type Buffer = Self; + + fn write(from: Self) -> Self::Buffer { + from + } + + fn read(from: Self::Buffer) -> Self { + from + } +} diff --git a/node-graph/gcore-shaders/src/shaders/buffer_struct/primitive.rs b/node-graph/gcore-shaders/src/shaders/buffer_struct/primitive.rs new file mode 100644 index 0000000000..84bc92d828 --- /dev/null +++ b/node-graph/gcore-shaders/src/shaders/buffer_struct/primitive.rs @@ -0,0 +1,135 @@ +use crate::shaders::buffer_struct::{BufferStruct, BufferStructIdentity}; +use bytemuck::Pod; +use core::marker::PhantomData; +use core::num::Wrapping; +use spirv_std::arch::IndexUnchecked; + +macro_rules! identity { + ($t:ty) => { + impl BufferStructIdentity for $t {} + }; +} + +identity!(()); +identity!(u8); +identity!(u16); +identity!(u32); +identity!(u64); +identity!(u128); +identity!(usize); +identity!(i8); +identity!(i16); +identity!(i32); +identity!(i64); +identity!(i128); +identity!(isize); +identity!(f32); +identity!(f64); + +identity!(spirv_std::arch::SubgroupMask); +identity!(spirv_std::memory::Semantics); +identity!(spirv_std::ray_tracing::RayFlags); +identity!(spirv_std::indirect_command::DrawIndirectCommand); +identity!(spirv_std::indirect_command::DrawIndexedIndirectCommand); +identity!(spirv_std::indirect_command::DispatchIndirectCommand); +identity!(spirv_std::indirect_command::DrawMeshTasksIndirectCommandEXT); +identity!(spirv_std::indirect_command::TraceRaysIndirectCommandKHR); +// not pod +// identity!(spirv_std::indirect_command::TraceRaysIndirectCommand2KHR); + +unsafe impl BufferStruct for bool { + type Buffer = u32; + + #[inline] + fn write(from: Self) -> Self::Buffer { + from as u32 + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + from != 0 + } +} + +unsafe impl BufferStruct for Wrapping +where + // unfortunately has to be Pod, even though AnyBitPattern would be sufficient, + // due to bytemuck doing `impl AnyBitPattern for T {}` + // see https://github.com/Lokathor/bytemuck/issues/164 + T::Buffer: Pod, +{ + type Buffer = Wrapping; + + #[inline] + fn write(from: Self) -> Self::Buffer { + Wrapping(T::write(from.0)) + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + Wrapping(T::read(from.0)) + } +} + +unsafe impl BufferStruct for PhantomData { + type Buffer = PhantomData; + + #[inline] + fn write(_: Self) -> Self::Buffer { + PhantomData {} + } + + #[inline] + fn read(_: Self::Buffer) -> Self { + PhantomData {} + } +} + +/// Potential problem: you can't impl this for an array of BufferStruct, as it'll conflict with this impl due to the +/// blanket impl on all BufferStructPlain types. +unsafe impl BufferStruct for [T; N] +where + // rust-gpu does not like `[T; N].map()` nor `core::array::from_fn()` nor transmuting arrays with a const generic + // length, so for now we need to require T: Default and T::Transfer: Default for all arrays. + T: Default, + // unfortunately has to be Pod, even though AnyBitPattern would be sufficient, + // due to bytemuck doing `impl AnyBitPattern for T {}` + // see https://github.com/Lokathor/bytemuck/issues/164 + T::Buffer: Pod + Default, +{ + type Buffer = [T::Buffer; N]; + + #[inline] + fn write(from: Self) -> Self::Buffer { + unsafe { + let mut ret = [T::Buffer::default(); N]; + for i in 0..N { + *ret.index_unchecked_mut(i) = T::write(*from.index_unchecked(i)); + } + ret + } + } + + #[inline] + fn read(from: Self::Buffer) -> Self { + unsafe { + let mut ret = [T::default(); N]; + for i in 0..N { + *ret.index_unchecked_mut(i) = T::read(*from.index_unchecked(i)); + } + ret + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_bool() { + for x in [false, true] { + assert_eq!(x, ::read(::write(x))); + } + } +} diff --git a/node-graph/gcore-shaders/src/shaders/mod.rs b/node-graph/gcore-shaders/src/shaders/mod.rs index 5e05a3b77e..1725a29516 100644 --- a/node-graph/gcore-shaders/src/shaders/mod.rs +++ b/node-graph/gcore-shaders/src/shaders/mod.rs @@ -1,5 +1,10 @@ +//! supporting infrastructure for shaders + +pub mod buffer_struct; + pub mod __private { pub use bytemuck; pub use glam; + pub use num_enum; pub use spirv_std; } diff --git a/node-graph/gcore/src/lib.rs b/node-graph/gcore/src/lib.rs index 3b098147ff..be1cbf967d 100644 --- a/node-graph/gcore/src/lib.rs +++ b/node-graph/gcore/src/lib.rs @@ -41,6 +41,7 @@ pub use graphene_core_shaders::AsU32; pub use graphene_core_shaders::blending; pub use graphene_core_shaders::choice_type; pub use graphene_core_shaders::color; +pub use graphene_core_shaders::shaders; pub use graphic::Graphic; pub use memo::MemoHash; pub use num_traits; diff --git a/node-graph/graster-nodes/Cargo.toml b/node-graph/graster-nodes/Cargo.toml index 7ed5357f00..c51c4485f9 100644 --- a/node-graph/graster-nodes/Cargo.toml +++ b/node-graph/graster-nodes/Cargo.toml @@ -48,6 +48,7 @@ bytemuck = { workspace = true } glam = { workspace = true } spirv-std = { workspace = true } num-traits = { workspace = true } +num_enum = { workspace = true } # Workspace std dependencies specta = { workspace = true, optional = true } diff --git a/node-graph/graster-nodes/src/adjustments.rs b/node-graph/graster-nodes/src/adjustments.rs index dc3bd15b82..48cbed4160 100644 --- a/node-graph/graster-nodes/src/adjustments.rs +++ b/node-graph/graster-nodes/src/adjustments.rs @@ -12,6 +12,8 @@ use graphene_core::table::Table; use graphene_core_shaders::color::Color; use graphene_core_shaders::context::Ctx; use graphene_core_shaders::registry::types::{AngleF32, PercentageF32, SignedPercentageF32}; +use node_macro::BufferStruct; +use num_enum::{FromPrimitive, IntoPrimitive}; #[cfg(not(feature = "std"))] use num_traits::float::Float; @@ -30,7 +32,7 @@ use num_traits::float::Float; // https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#:~:text=%27clrL%27%20%3D%20Color%20Lookup // https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#:~:text=Color%20Lookup%20(Photoshop%20CS6 -#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, node_macro::ChoiceType, bytemuck::NoUninit)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, node_macro::ChoiceType, bytemuck::NoUninit, BufferStruct, FromPrimitive, IntoPrimitive)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[widget(Dropdown)] #[repr(u32)] @@ -549,7 +551,8 @@ fn vibrance>( } /// Color Channel -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType)] +#[repr(u32)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType, BufferStruct, FromPrimitive, IntoPrimitive)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[widget(Radio)] pub enum RedGreenBlue { @@ -560,7 +563,7 @@ pub enum RedGreenBlue { } /// Color Channel -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType, bytemuck::NoUninit)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType, bytemuck::NoUninit, BufferStruct, FromPrimitive, IntoPrimitive)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[widget(Radio)] #[repr(u32)] diff --git a/node-graph/node-macro/src/buffer_struct.rs b/node-graph/node-macro/src/buffer_struct.rs new file mode 100644 index 0000000000..1caa94b997 --- /dev/null +++ b/node-graph/node-macro/src/buffer_struct.rs @@ -0,0 +1,261 @@ +use crate::crate_ident::CrateIdent; +use proc_macro2::{Ident, Span, TokenStream}; +use quote::{ToTokens, format_ident, quote}; +use std::collections::HashSet; +use syn::punctuated::Punctuated; +use syn::visit_mut::VisitMut; +use syn::{Fields, GenericParam, Generics, Item, ItemEnum, ItemStruct, Meta, MetaList, Path, PathSegment, Result, Token, TypeParam, TypeParamBound, visit_mut}; + +pub fn derive_buffer_struct(crate_ident: &CrateIdent, content: proc_macro::TokenStream) -> Result { + let item = syn::parse::(content)?; + match &item { + Item::Enum(item) => derive_buffer_struct_enum(crate_ident, item), + Item::Struct(item) => derive_buffer_struct_struct(crate_ident, item), + _ => Err(syn::Error::new_spanned(&item, "Expected a struct or an enum")), + } +} + +pub fn derive_buffer_struct_enum(crate_ident: &CrateIdent, item: &ItemEnum) -> Result { + let gcore_shaders = crate_ident.gcore_shaders()?; + let mod_buffer_struct = quote!(#gcore_shaders::shaders::buffer_struct); + let reexport = quote!(#gcore_shaders::shaders::__private); + + if !item.generics.params.is_empty() { + return Err(syn::Error::new_spanned(&item.generics, "enum must not have any generics")); + } + + let enum_requirements_error = || { + syn::Error::new( + Span::call_site(), + "deriving `BufferStruct` on an enum requires `#[repr(u32)]` and `#[derive(num_enum::FromPrimitive, num_enum::IntoPrimitive)]`", + ) + }; + let repr_path = Path::from(format_ident!("repr")); + let repr = item + .attrs + .iter() + .filter_map(|a| match &a.meta { + Meta::List(MetaList { path, tokens, .. }) if *path == repr_path => Some(tokens), + _ => None, + }) + .next() + .ok_or_else(enum_requirements_error)?; + + let ident = &item.ident; + Ok(quote! { + unsafe impl #mod_buffer_struct::BufferStruct for #ident + { + type Buffer = #repr; + + fn write(from: Self) -> Self::Buffer { + <#repr as From>::from(from) + } + + fn read(from: Self::Buffer) -> Self { + ::from_primitive(from) + } + } + }) +} + +/// see `BufferStruct` docs +/// +/// This is also largely copied from my (@firestar99) project and adjusted +pub fn derive_buffer_struct_struct(crate_ident: &CrateIdent, item: &ItemStruct) -> Result { + let gcore_shaders = crate_ident.gcore_shaders()?; + let mod_buffer_struct = quote!(#gcore_shaders::shaders::buffer_struct); + let reexport = quote!(#gcore_shaders::shaders::__private); + + let generics = item + .generics + .params + .iter() + .filter_map(|g| match g { + GenericParam::Lifetime(_) => None, + GenericParam::Type(t) => Some(t.ident.clone()), + GenericParam::Const(c) => Some(c.ident.clone()), + }) + .collect(); + + let mut members_buffer = Punctuated::::new(); + let mut write = Punctuated::::new(); + let mut read = Punctuated::::new(); + let mut gen_name_gen = GenericNameGen::new(); + let mut gen_ref_tys = Vec::new(); + let (members_buffer, write, read) = match &item.fields { + Fields::Named(named) => { + for f in &named.named { + let name = f.ident.as_ref().unwrap(); + let mut ty = f.ty.clone(); + let mut visitor = GenericsVisitor::new(&item.ident, &generics); + visit_mut::visit_type_mut(&mut visitor, &mut ty); + if visitor.found_generics { + gen_ref_tys.push(f.ty.clone()); + let gen_ident = gen_name_gen.next(); + members_buffer.push(quote!(#name: #gen_ident)); + } else { + members_buffer.push(quote! { + #name: <#ty as #mod_buffer_struct::BufferStruct>::Buffer + }); + } + + write.push(quote! { + #name: <#ty as #mod_buffer_struct::BufferStruct>::write(from.#name) + }); + read.push(quote! { + #name: <#ty as #mod_buffer_struct::BufferStruct>::read(from.#name) + }); + } + (quote!({#members_buffer}), quote!(Self::Buffer {#write}), quote!(Self {#read})) + } + Fields::Unnamed(unnamed) => { + for (i, f) in unnamed.unnamed.iter().enumerate() { + let mut ty = f.ty.clone(); + let mut visitor = GenericsVisitor::new(&item.ident, &generics); + visit_mut::visit_type_mut(&mut visitor, &mut ty); + if visitor.found_generics { + gen_ref_tys.push(f.ty.clone()); + members_buffer.push(gen_name_gen.next().into_token_stream()); + } else { + members_buffer.push(quote! { + <#ty as #mod_buffer_struct::BufferStruct>::Buffer + }); + } + + let index = syn::Index::from(i); + write.push(quote! { + <#ty as #mod_buffer_struct::BufferStruct>::write(from.#index) + }); + read.push(quote! { + <#ty as #mod_buffer_struct::BufferStruct>::read(from.#index) + }); + } + (quote!((#members_buffer);), quote!(Self::Buffer(#write)), quote!(Self(#read))) + } + Fields::Unit => (quote!(;), quote!(let _ = from; Self::Buffer {}), quote!(let _ = from; Self::Shader {})), + }; + + let generics_decl = &item.generics; + let generics_ref = decl_to_ref(item.generics.params.iter()); + let generics_where = gen_ref_tys + .iter() + .map(|ty| quote!(#ty: #mod_buffer_struct::BufferStruct)) + .collect::>() + .into_token_stream(); + + let generics_decl_any = gen_name_gen.decl(quote! { + #reexport::bytemuck::Pod + Send + Sync + }); + let generics_ref_buffer = gen_ref_tys + .iter() + .map(|ty| quote!(<#ty as #mod_buffer_struct::BufferStruct>::Buffer)) + .collect::>() + .into_token_stream(); + + let vis = &item.vis; + let ident = &item.ident; + let buffer_ident = format_ident!("{}Buffer", ident); + Ok(quote! { + #[repr(C)] + #[derive(Copy, Clone, #reexport::bytemuck::Zeroable, #reexport::bytemuck::Pod)] + #vis struct #buffer_ident #generics_decl_any #members_buffer + + unsafe impl #generics_decl #mod_buffer_struct::BufferStruct for #ident #generics_ref + where + #ident #generics_ref: Copy, + #generics_where + { + type Buffer = #buffer_ident <#generics_ref_buffer>; + + fn write(from: Self) -> Self::Buffer { + #write + } + + fn read(from: Self::Buffer) -> Self { + #read + } + } + }) +} + +struct GenericsVisitor<'a> { + self_ident: &'a Ident, + generics: &'a HashSet, + found_generics: bool, +} + +impl<'a> GenericsVisitor<'a> { + pub fn new(self_ident: &'a Ident, generics: &'a HashSet) -> Self { + Self { + self_ident, + generics, + found_generics: false, + } + } +} + +impl VisitMut for GenericsVisitor<'_> { + fn visit_ident_mut(&mut self, i: &mut Ident) { + if self.generics.contains(i) { + self.found_generics = true; + } + visit_mut::visit_ident_mut(self, i); + } + + fn visit_path_segment_mut(&mut self, i: &mut PathSegment) { + if i.ident.to_string() == "Self" { + i.ident = self.self_ident.clone(); + } + visit_mut::visit_path_segment_mut(self, i); + } +} + +struct GenericNameGen(u32); + +impl GenericNameGen { + pub fn new() -> Self { + Self(0) + } + + pub fn next(&mut self) -> Ident { + let i = self.0; + self.0 += 1; + format_ident!("T{}", i) + } + + pub fn decl(self, ty: TokenStream) -> Generics { + let params: Punctuated = (0..self.0) + .map(|i| { + GenericParam::Type(TypeParam { + attrs: Vec::new(), + ident: format_ident!("T{}", i), + colon_token: Some(Default::default()), + bounds: Punctuated::from_iter([TypeParamBound::Verbatim(ty.clone())]), + eq_token: None, + default: None, + }) + }) + .collect(); + if !params.is_empty() { + Generics { + lt_token: Some(Default::default()), + params, + gt_token: Some(Default::default()), + where_clause: None, + } + } else { + Generics::default() + } + } +} + +fn decl_to_ref<'a>(generics: impl Iterator) -> TokenStream { + let out = generics + .map(|generic| match generic { + GenericParam::Lifetime(l) => l.lifetime.to_token_stream(), + GenericParam::Type(t) => t.ident.to_token_stream(), + GenericParam::Const(c) => c.ident.to_token_stream(), + }) + .collect::>(); + if out.is_empty() { TokenStream::new() } else { quote!(<#out>) } +} diff --git a/node-graph/node-macro/src/lib.rs b/node-graph/node-macro/src/lib.rs index 3ff09d3753..53901d153d 100644 --- a/node-graph/node-macro/src/lib.rs +++ b/node-graph/node-macro/src/lib.rs @@ -1,7 +1,9 @@ +use crate::crate_ident::CrateIdent; use proc_macro::TokenStream; use proc_macro_error2::proc_macro_error; use syn::GenericParam; +mod buffer_struct; mod codegen; mod crate_ident; mod derive_choice_type; @@ -30,3 +32,10 @@ pub fn node(attr: TokenStream, item: TokenStream) -> TokenStream { pub fn derive_choice_type(input_item: TokenStream) -> TokenStream { derive_choice_type::derive_choice_type_impl(input_item.into()).unwrap_or_else(|err| err.to_compile_error()).into() } + +/// Derive a struct to implement `ShaderStruct`, see that for docs. +#[proc_macro_derive(BufferStruct)] +pub fn derive_buffer_struct(input_item: TokenStream) -> TokenStream { + let crate_ident = CrateIdent::default(); + TokenStream::from(buffer_struct::derive_buffer_struct(&crate_ident, input_item.into()).unwrap_or_else(|err| err.to_compile_error())) +} diff --git a/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs b/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs index 1f9e2e37e0..01559eaff7 100644 --- a/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs +++ b/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs @@ -120,6 +120,16 @@ impl PerPixelAdjustCodegen<'_> { ParamType::Uniform => Some(quote! {#ident: #ty}), }) .collect::>(); + let uniform_struct_ident = &self.uniform_struct_ident; + let uniform_struct = parse_quote! { + #[repr(C)] + #[derive(Copy, Clone)] + pub struct #uniform_struct_ident { + #(pub #uniform_members),* + } + }; + let uniform_struct_shader_struct_derive = crate::buffer_struct::derive_buffer_struct_struct(&self.crate_ident, &uniform_struct)?; + let image_params = self .params .iter() @@ -140,7 +150,6 @@ impl PerPixelAdjustCodegen<'_> { let entry_point_mod = &self.entry_point_mod; let entry_point_name = &self.entry_point_name_ident; - let uniform_struct_ident = &self.uniform_struct_ident; Ok(quote! { pub mod #entry_point_mod { use super::*; @@ -152,19 +161,17 @@ impl PerPixelAdjustCodegen<'_> { pub const #entry_point_name: &str = core::concat!(core::module_path!(), "::entry_point"); - #[repr(C)] - #[derive(Copy, Clone, #reexport::bytemuck::NoUninit)] - pub struct #uniform_struct_ident { - #(pub #uniform_members),* - } + #uniform_struct + #uniform_struct_shader_struct_derive #[spirv(fragment)] pub fn entry_point( #[spirv(frag_coord)] frag_coord: Vec4, color_out: &mut Vec4, - #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] uniform: &Uniform, + #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] uniform: &UniformBuffer, #(#image_params),* ) { + let uniform = ::read(*uniform); let texel_coord = frag_coord.xy().as_uvec2(); let color: Color = #fn_name(#context, #(#call_args),*); *color_out = color.to_vec4(); diff --git a/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs b/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs index 928f35a0b0..aa567a5b93 100644 --- a/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs +++ b/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs @@ -1,8 +1,8 @@ use crate::Context; use crate::shader_runtime::{FULLSCREEN_VERTEX_SHADER_NAME, ShaderRuntime}; -use bytemuck::NoUninit; use futures::lock::Mutex; use graphene_core::raster_types::{GPU, Raster}; +use graphene_core::shaders::buffer_struct::BufferStruct; use graphene_core::table::{Table, TableRow}; use std::borrow::Cow; use std::collections::HashMap; @@ -27,7 +27,7 @@ impl PerPixelAdjustShaderRuntime { } impl ShaderRuntime { - pub async fn run_per_pixel_adjust(&self, shaders: &Shaders<'_>, textures: Table>, args: Option<&T>) -> Table> { + pub async fn run_per_pixel_adjust(&self, shaders: &Shaders<'_>, textures: Table>, args: Option<&T>) -> Table> { let mut cache = self.per_pixel_adjust.pipeline_cache.lock().await; let pipeline = cache .entry(shaders.fragment_shader_name.to_owned()) @@ -38,7 +38,7 @@ impl ShaderRuntime { device.create_buffer_init(&BufferInitDescriptor { label: Some(&format!("{} arg buffer", pipeline.name.as_str())), usage: BufferUsages::STORAGE, - contents: bytemuck::bytes_of(args), + contents: bytemuck::bytes_of(&T::write(*args)), }) }); pipeline.dispatch(&self.context, textures, arg_buffer)