shaders: codegen for per_pixel_adjust shader nodes

Author: Firestar99

node-graph/gcore-shaders/src/color/color_types.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(not(feature = "std"))]
 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

@@ -52,6 +52,7 @@ fn luminance<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	luminance_calc: LuminanceCalculation,
 ) -> T {
@@ -77,6 +78,7 @@ fn gamma_correction<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	#[default(2.2)]
 	#[range((0.01, 10.))]
@@ -98,6 +100,7 @@ fn extract_channel<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	channel: RedGreenBlueAlpha,
 ) -> T {
@@ -122,6 +125,7 @@ fn make_opaque<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 ) -> T {
 	input.adjust(|color| {
@@ -148,6 +152,7 @@ fn brightness_contrast<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	brightness: SignedPercentageF32,
 	contrast: SignedPercentageF32,
@@ -238,6 +243,7 @@ fn levels<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(0.)] shadows: PercentageF32,
 	#[default(50.)] midtones: PercentageF32,
@@ -306,6 +312,7 @@ fn black_and_white<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(Color::BLACK)] tint: Color,
 	#[default(40.)]
@@ -379,6 +386,7 @@ fn hue_saturation<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	hue_shift: AngleF32,
 	saturation_shift: SignedPercentageF32,
@@ -414,6 +422,7 @@ fn invert<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 ) -> T {
 	input.adjust(|color| {
@@ -437,6 +446,7 @@ fn threshold<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	#[default(50.)] min_luminance: PercentageF32,
 	#[default(100.)] max_luminance: PercentageF32,
@@ -483,6 +493,7 @@ fn vibrance<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 	vibrance: SignedPercentageF32,
 ) -> T {
@@ -649,6 +660,7 @@ fn channel_mixer<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 
 	monochrome: bool,
@@ -778,6 +790,7 @@ fn selective_color<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut image: T,
 
 	mode: RelativeAbsolute,
@@ -921,6 +934,7 @@ fn posterize<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	#[default(4)]
 	#[hard_min(2.)]
@@ -955,6 +969,7 @@ fn exposure<T: Adjust<Color>>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	mut input: T,
 	exposure: f32,
 	offset: f32,

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

@@ -141,6 +141,7 @@ fn blend<T: Blend<Color> + Send>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	over: T,
 	#[expose]
 	#[implementations(
@@ -149,6 +150,7 @@ fn blend<T: Blend<Color> + Send>(
 		Table<GradientStops>,
 		GradientStops,
 	)]
+	#[gpu_image]
 	under: T,
 	blend_mode: BlendMode,
 	#[default(100.)] opacity: PercentageF32,
@@ -165,6 +167,7 @@ fn color_overlay<T: Adjust<Color>>(
 		Table<GradientStops>,
 		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;
@@ -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.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),
+		}
+	}
+}