shaders: codegen for per_pixel_adjust shader nodes

Author: Firestar99

node-graph/gcore-shaders/src/color/color.rs

@@ -3,6 +3,7 @@ use super::discrete_srgb::{float_to_srgb_u8, srgb_u8_to_float};
 use bytemuck::{Pod, Zeroable};
 use core::fmt::Debug;
 use core::hash::Hash;
+use glam::Vec4;
 use half::f16;
 #[cfg(target_arch = "spirv")]
 use num_traits::Euclid;
@@ -1075,6 +1076,21 @@ impl Color {
 			..*self
 		}
 	}
+
+	#[inline(always)]
+	pub const fn from_vec4(vec: Vec4) -> Self {
+		Self {
+			red: vec.x,
+			green: vec.y,
+			blue: vec.z,
+			alpha: vec.w,
+		}
+	}
+
+	#[inline(always)]
+	pub fn to_vec4(&self) -> Vec4 {
+		Vec4::new(self.red, self.green, self.blue, self.alpha)
+	}
 }
 
 #[cfg(test)]

node-graph/graster-nodes/src/adjustments.rs

@@ -50,6 +50,7 @@ fn luminance<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	luminance_calc: LuminanceCalculation,
 ) -> T {
@@ -74,6 +75,7 @@ fn gamma_correction<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	#[default(2.2)]
 	#[range((0.01, 10.))]
@@ -94,6 +96,7 @@ fn extract_channel<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	channel: RedGreenBlueAlpha,
 ) -> T {
@@ -117,6 +120,7 @@ fn make_opaque<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 ) -> T {
 	input.adjust(|color| {
@@ -142,6 +146,7 @@ fn brightness_contrast<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	brightness: SignedPercentage,
 	contrast: SignedPercentage,
@@ -231,6 +236,7 @@ fn levels<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(0.)] shadows: Percentage,
 	#[default(50.)] midtones: Percentage,
@@ -298,6 +304,7 @@ async fn black_and_white<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(Color::BLACK)] tint: Color,
 	#[default(40.)]
@@ -370,6 +377,7 @@ async fn hue_saturation<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	hue_shift: Angle,
 	saturation_shift: SignedPercentage,
@@ -404,6 +412,7 @@ async fn invert<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 ) -> T {
 	input.adjust(|color| {
@@ -426,6 +435,7 @@ async fn threshold<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(50.)] min_luminance: Percentage,
 	#[default(100.)] max_luminance: Percentage,
@@ -471,6 +481,7 @@ async fn vibrance<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	vibrance: SignedPercentage,
 ) -> T {
@@ -636,6 +647,7 @@ async fn channel_mixer<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 
 	monochrome: bool,
@@ -764,6 +776,7 @@ async fn selective_color<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 
 	mode: RelativeAbsolute,
@@ -906,6 +919,7 @@ async fn posterize<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	#[default(4)]
 	#[hard_min(2.)]
@@ -939,6 +953,7 @@ async fn exposure<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	exposure: f64,
 	offset: f64,

node-graph/graster-nodes/src/blending_nodes.rs

@@ -130,13 +130,15 @@ async fn blend<T: Blend<Color> + Send>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	over: T,
 	#[expose]
 	#[implementations(
 		Color,
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	under: T,
 	blend_mode: BlendMode,
 	#[default(100.)] opacity: Percentage,
@@ -152,6 +154,7 @@ fn color_overlay<T: Adjust<Color>>(
 		Table<Raster<CPU>>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(Color::BLACK)] color: Color,
 	blend_mode: BlendMode,

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;
@@ -347,6 +347,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]
@@ -435,6 +436,8 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result<TokenStre
 				NODE_METADATA.lock().unwrap().insert(#identifier(), metadata);
 			}
 		}
+
+		#shader_entry_point
 	})
 }
 
@@ -637,6 +640,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

@@ -104,6 +104,8 @@ impl Parse for ParsedWidgetOverride {
 
 #[derive(Debug)]
 pub(crate) enum ParsedField {
+	/// a param of any kind, either a concrete type or a generic type with a set of possible types specified via
+	/// `#[implementation(type)]`
 	Regular {
 		pat_ident: PatIdent,
 		name: Option<LitStr>,
@@ -121,7 +123,9 @@ pub(crate) enum ParsedField {
 		number_step: Option<LitFloat>,
 		implementations: Punctuated<Type, Comma>,
 		unit: Option<LitStr>,
+		gpu_image: bool,
 	},
+	/// a param of `impl Node` with `#[implementation(in -> out)]`
 	Node {
 		pat_ident: PatIdent,
 		name: Option<LitStr>,
@@ -533,6 +537,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 '{}': {}\nUSAGE EXAMPLE: #[step(2.)]", ident, e)))
 		})
 		.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
@@ -597,6 +602,7 @@ fn parse_field(pat_ident: PatIdent, ty: Type, attrs: &[Attribute]) -> syn::Resul
 			value_source,
 			implementations,
 			unit,
+			gpu_image,
 		})
 	}
 }

node-graph/node-macro/src/shader_nodes.rs renamed to 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),
+		}
+	}
+}