shaders: codegen for per_pixel_adjust shader nodes

node-graph/node-macro/src/codegen.rs

@@ -1,7 +1,7 @@
 use crate::parsing::*;
 use convert_case::{Case, Casing};
 use proc_macro_crate::FoundCrate;
-use proc_macro2::TokenStream as TokenStream2;
+use proc_macro2::{TokenStream as TokenStream2, TokenStream};
 use quote::{ToTokens, format_ident, quote, quote_spanned};
 use std::sync::atomic::AtomicU64;
 use syn::punctuated::Punctuated;
@@ -295,6 +295,7 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result<TokenStre
 
 	let cfg = crate::shader_nodes::modify_cfg(attributes);
 	let node_input_accessor = generate_node_input_references(parsed, fn_generics, &field_idents, &graphene_core, &identifier, &cfg);
+	let shader_entry_point = attributes.shader_node.as_ref().map(|n| n.codegen_shader_entry_point(parsed)).unwrap_or(Ok(TokenStream::new()))?;
 	Ok(quote! {
 		/// Underlying implementation for [#struct_name]
 		#[inline]
@@ -384,6 +385,8 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result<TokenStre
 				NODE_METADATA.lock().unwrap().insert(#identifier(), metadata);
 			}
 		}
+
+		#shader_entry_point
 	})
 }
 
@@ -586,6 +589,7 @@ fn generate_register_node_impl(parsed: &ParsedNodeFn, field_names: &[&Ident], st
 	})
 }
 
+use crate::shader_nodes::CodegenShaderEntryPoint;
 use syn::visit_mut::VisitMut;
 use syn::{GenericArgument, Lifetime, Type};
 

node-graph/node-macro/src/parsing.rs

@@ -120,6 +120,8 @@ pub enum ParsedFieldType {
 	Node(NodeParsedField),
 }
 
+/// a param of any kind, either a concrete type or a generic type with a set of possible types specified via
+/// `#[implementation(type)]`
 #[derive(Clone, Debug)]
 pub struct RegularParsedField {
 	pub ty: Type,
@@ -131,8 +133,10 @@ pub struct RegularParsedField {
 	pub number_hard_max: Option<LitFloat>,
 	pub number_mode_range: Option<ExprTuple>,
 	pub implementations: Punctuated<Type, Comma>,
+	pub gpu_image: bool,
 }
 
+/// a param of `impl Node` with `#[implementation(in -> out)]`
 #[derive(Clone, Debug)]
 pub struct NodeParsedField {
 	pub input_type: Type,
@@ -529,6 +533,7 @@ fn parse_field(pat_ident: PatIdent, ty: Type, attrs: &[Attribute]) -> syn::Resul
 				.map_err(|e| Error::new_spanned(attr, format!("Invalid `step` for argument '{ident}': {e}\nUSAGE EXAMPLE: #[step(2.)]")))
 		})
 		.transpose()?;
+	let gpu_image = extract_attribute(attrs, "gpu_image").is_some();
 
 	let (is_node, node_input_type, node_output_type) = parse_node_type(&ty);
 	let description = attrs
@@ -590,6 +595,7 @@ fn parse_field(pat_ident: PatIdent, ty: Type, attrs: &[Attribute]) -> syn::Resul
 				ty,
 				value_source,
 				implementations,
+				gpu_image,
 			}),
 			name,
 			description,

node-graph/node-macro/src/shader_nodes/mod.rs

@@ -1,10 +1,13 @@
-use crate::parsing::NodeFnAttributes;
+use crate::parsing::{NodeFnAttributes, ParsedNodeFn};
+use crate::shader_nodes::per_pixel_adjust::PerPixelAdjust;
 use proc_macro2::{Ident, TokenStream};
 use quote::quote;
-use strum::{EnumString, VariantNames};
+use strum::VariantNames;
 use syn::Error;
 use syn::parse::{Parse, ParseStream};
 
+pub mod per_pixel_adjust;
+
 pub const STD_FEATURE_GATE: &str = "std";
 
 pub fn modify_cfg(attributes: &NodeFnAttributes) -> TokenStream {
@@ -16,17 +19,33 @@ pub fn modify_cfg(attributes: &NodeFnAttributes) -> TokenStream {
 	}
 }
 
-#[derive(Debug, EnumString, VariantNames)]
+#[derive(Debug, VariantNames)]
 pub(crate) enum ShaderNodeType {
-	PerPixelAdjust,
+	PerPixelAdjust(PerPixelAdjust),
 }
 
 impl Parse for ShaderNodeType {
 	fn parse(input: ParseStream) -> syn::Result<Self> {
 		let ident: Ident = input.parse()?;
 		Ok(match ident.to_string().as_str() {
-			"PerPixelAdjust" => ShaderNodeType::PerPixelAdjust,
+			"PerPixelAdjust" => ShaderNodeType::PerPixelAdjust(PerPixelAdjust::parse(input)?),
 			_ => return Err(Error::new_spanned(&ident, format!("attr 'shader_node' must be one of {:?}", Self::VARIANTS))),
 		})
 	}
 }
+
+pub trait CodegenShaderEntryPoint {
+	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream>;
+}
+
+impl CodegenShaderEntryPoint for ShaderNodeType {
+	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
+		if parsed.is_async {
+			return Err(Error::new_spanned(&parsed.fn_name, "Shader nodes must not be async"));
+		}
+
+		match self {
+			ShaderNodeType::PerPixelAdjust(x) => x.codegen_shader_entry_point(parsed),
+		}
+	}
+}

node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs

@@ -0,0 +1,110 @@
+use crate::parsing::{ParsedFieldType, ParsedNodeFn, RegularParsedField};
+use crate::shader_nodes::CodegenShaderEntryPoint;
+use proc_macro2::{Ident, TokenStream};
+use quote::{ToTokens, format_ident, quote};
+use std::borrow::Cow;
+use syn::parse::{Parse, ParseStream};
+
+#[derive(Debug)]
+pub struct PerPixelAdjust {}
+
+impl Parse for PerPixelAdjust {
+	fn parse(_input: ParseStream) -> syn::Result<Self> {
+		Ok(Self {})
+	}
+}
+
+impl CodegenShaderEntryPoint for PerPixelAdjust {
+	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
+		let fn_name = &parsed.fn_name;
+		let gpu_mod = format_ident!("{}_gpu", parsed.fn_name);
+		let spirv_image_ty = quote!(Image2d);
+
+		// bindings for images start at 1
+		let mut binding_cnt = 0;
+		let params = parsed
+			.fields
+			.iter()
+			.map(|f| {
+				let ident = &f.pat_ident;
+				match &f.ty {
+					ParsedFieldType::Node { .. } => Err(syn::Error::new_spanned(ident, "PerPixelAdjust shader nodes cannot accept other nodes as generics")),
+					ParsedFieldType::Regular(RegularParsedField { gpu_image: false, ty, .. }) => Ok(Param {
+						ident: Cow::Borrowed(&ident.ident),
+						ty: Cow::Owned(ty.to_token_stream()),
+						param_type: ParamType::Uniform,
+					}),
+					ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. }) => {
+						binding_cnt += 1;
+						Ok(Param {
+							ident: Cow::Owned(format_ident!("image_{}", &ident.ident)),
+							ty: Cow::Borrowed(&spirv_image_ty),
+							param_type: ParamType::Image { binding: binding_cnt },
+						})
+					}
+				}
+			})
+			.collect::<syn::Result<Vec<_>>>()?;
+
+		let uniform_members = params
+			.iter()
+			.filter_map(|Param { ident, ty, param_type }| match param_type {
+				ParamType::Image { .. } => None,
+				ParamType::Uniform => Some(quote! {#ident: #ty}),
+			})
+			.collect::<Vec<_>>();
+		let image_params = params
+			.iter()
+			.filter_map(|Param { ident, ty, param_type }| match param_type {
+				ParamType::Image { binding } => Some(quote! {#[spirv(descriptor_set = 0, binding = #binding)] #ident: &#ty}),
+				ParamType::Uniform => None,
+			})
+			.collect::<Vec<_>>();
+		let call_args = params
+			.iter()
+			.map(|Param { ident, param_type, .. }| match param_type {
+				ParamType::Image { .. } => quote!(Color::from_vec4(#ident.fetch_with(texel_coord, lod(0)))),
+				ParamType::Uniform => quote!(uniform.#ident),
+			})
+			.collect::<Vec<_>>();
+		let context = quote!(());
+
+		Ok(quote! {
+			pub mod #gpu_mod {
+				use super::*;
+				use graphene_core_shaders::color::Color;
+				use spirv_std::spirv;
+				use spirv_std::glam::{Vec4, Vec4Swizzles};
+				use spirv_std::image::{Image2d, ImageWithMethods};
+				use spirv_std::image::sample_with::lod;
+
+				pub struct Uniform {
+					#(#uniform_members),*
+				}
+
+				#[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,
+					#(#image_params),*
+				) {
+					let texel_coord = frag_coord.xy().as_uvec2();
+					let color: Color = #fn_name(#context, #(#call_args),*);
+					*color_out = color.to_vec4();
+				}
+			}
+		})
+	}
+}
+
+struct Param<'a> {
+	ident: Cow<'a, Ident>,
+	ty: Cow<'a, TokenStream>,
+	param_type: ParamType,
+}
+
+enum ParamType {
+	Image { binding: u32 },
+	Uniform,
+}