shader-rt: properly cfg out the gpu node

Author: Firestar99

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, TokenStream};
+use proc_macro2::TokenStream as TokenStream2;
 use quote::{ToTokens, format_ident, quote, quote_spanned};
 use std::sync::atomic::AtomicU64;
 use syn::punctuated::Punctuated;
@@ -295,11 +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, shader_gpu_node) = attributes
-		.shader_node
-		.as_ref()
-		.map::<syn::Result<_>, _>(|n| Ok((n.codegen_shader_entry_point(parsed)?, n.codegen_gpu_node(parsed)?)))
-		.unwrap_or(Ok((TokenStream::new(), TokenStream::new())))?;
+	let ShaderTokens { shader_entry_point, gpu_node } = attributes.shader_node.as_ref().map(|n| n.codegen(parsed, &cfg)).unwrap_or(Ok(ShaderTokens::default()))?;
 
 	Ok(quote! {
 		/// Underlying implementation for [#struct_name]
@@ -392,7 +388,7 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result<TokenStre
 
 		#shader_entry_point
 
-		#shader_gpu_node
+		#gpu_node
 	})
 }
 
@@ -595,7 +591,7 @@ fn generate_register_node_impl(parsed: &ParsedNodeFn, field_names: &[&Ident], st
 	})
 }
 
-use crate::shader_nodes::CodegenShaderEntryPoint;
+use crate::shader_nodes::{ShaderCodegen, ShaderTokens};
 use syn::visit_mut::VisitMut;
 use syn::{GenericArgument, Lifetime, Type};
 

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

@@ -38,13 +38,12 @@ impl Parse for ShaderNodeType {
 	}
 }
 
-pub trait CodegenShaderEntryPoint {
-	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream>;
-	fn codegen_gpu_node(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream>;
+pub trait ShaderCodegen {
+	fn codegen(&self, parsed: &ParsedNodeFn, node_cfg: &TokenStream) -> syn::Result<ShaderTokens>;
 }
 
-impl CodegenShaderEntryPoint for ShaderNodeType {
-	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
+impl ShaderCodegen for ShaderNodeType {
+	fn codegen(&self, parsed: &ParsedNodeFn, node_cfg: &TokenStream) -> syn::Result<ShaderTokens> {
 		match self {
 			ShaderNodeType::GpuNode => (),
 			_ => {
@@ -55,15 +54,14 @@ impl CodegenShaderEntryPoint for ShaderNodeType {
 		}
 
 		match self {
-			ShaderNodeType::GpuNode => Ok(TokenStream::new()),
-			ShaderNodeType::PerPixelAdjust(x) => x.codegen_shader_entry_point(parsed),
+			ShaderNodeType::GpuNode => Ok(ShaderTokens::default()),
+			ShaderNodeType::PerPixelAdjust(x) => x.codegen(parsed, node_cfg),
 		}
 	}
+}
 
-	fn codegen_gpu_node(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
-		match self {
-			ShaderNodeType::GpuNode => Ok(TokenStream::new()),
-			ShaderNodeType::PerPixelAdjust(x) => x.codegen_gpu_node(parsed),
-		}
-	}
+#[derive(Clone, Default)]
+pub struct ShaderTokens {
+	pub shader_entry_point: TokenStream,
+	pub gpu_node: TokenStream,
 }

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

@@ -1,13 +1,13 @@
 use crate::parsing::{Input, NodeFnAttributes, ParsedField, ParsedFieldType, ParsedNodeFn, RegularParsedField};
-use crate::shader_nodes::{CodegenShaderEntryPoint, ShaderNodeType};
+use crate::shader_nodes::{ShaderCodegen, ShaderNodeType, ShaderTokens};
 use convert_case::{Case, Casing};
 use proc_macro_crate::FoundCrate;
 use proc_macro2::{Ident, Span, TokenStream};
 use quote::{ToTokens, format_ident, quote};
 use std::borrow::Cow;
 use syn::parse::{Parse, ParseStream};
 use syn::punctuated::Punctuated;
-use syn::{Path, Token, TraitBound, TraitBoundModifier, Type, TypeImplTrait, TypeParamBound};
+use syn::{Token, TraitBound, TraitBoundModifier, Type, TypeImplTrait, TypeParamBound};
 
 #[derive(Debug, Clone)]
 pub struct PerPixelAdjust {}
@@ -18,10 +18,19 @@ impl Parse for PerPixelAdjust {
 	}
 }
 
-impl CodegenShaderEntryPoint for PerPixelAdjust {
+impl ShaderCodegen for PerPixelAdjust {
+	fn codegen(&self, parsed: &ParsedNodeFn, node_cfg: &TokenStream) -> syn::Result<ShaderTokens> {
+		Ok(ShaderTokens {
+			shader_entry_point: self.codegen_shader_entry_point(parsed)?,
+			gpu_node: self.codegen_gpu_node(parsed, node_cfg)?,
+		})
+	}
+}
+
+impl PerPixelAdjust {
 	fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
 		let fn_name = &parsed.fn_name;
-		let gpu_mod = format_ident!("{}_gpu_entry_point", parsed.fn_name);
+		let gpu_mod = format_ident!("{}_gpu_entry_point", fn_name);
 		let spirv_image_ty = quote!(Image2d);
 
 		// bindings for images start at 1
@@ -101,7 +110,7 @@ impl CodegenShaderEntryPoint for PerPixelAdjust {
 		})
 	}
 
-	fn codegen_gpu_node(&self, parsed: &ParsedNodeFn) -> syn::Result<TokenStream> {
+	fn codegen_gpu_node(&self, parsed: &ParsedNodeFn, node_cfg: &TokenStream) -> syn::Result<TokenStream> {
 		let fn_name = format_ident!("{}_gpu", parsed.fn_name);
 		let struct_name = format_ident!("{}", fn_name.to_string().to_case(Case::Pascal));
 		let mod_name = fn_name.clone();
@@ -127,21 +136,22 @@ impl CodegenShaderEntryPoint for PerPixelAdjust {
 				ParsedFieldType::Node { .. } => Err(syn::Error::new_spanned(&f.pat_ident, "PerPixelAdjust shader nodes cannot accept other nodes as generics")),
 			})
 			.collect::<syn::Result<_>>()?;
+
 		let body = quote! {
 			{
 
 			}
 		};
 
-		crate::codegen::generate_node_code(&ParsedNodeFn {
+		let gpu_node = crate::codegen::generate_node_code(&ParsedNodeFn {
 			vis: parsed.vis.clone(),
 			attributes: NodeFnAttributes {
 				shader_node: Some(ShaderNodeType::GpuNode),
 				..parsed.attributes.clone()
 			},
 			fn_name,
 			struct_name,
-			mod_name,
+			mod_name: mod_name.clone(),
 			fn_generics: vec![],
 			where_clause: None,
 			input: Input {
@@ -152,7 +162,7 @@ impl CodegenShaderEntryPoint for PerPixelAdjust {
 						paren_token: None,
 						modifier: TraitBoundModifier::None,
 						lifetimes: None,
-						path: Path::from(format_ident!("Ctx")),
+						path: syn::parse2(quote!(#gcore::context::Ctx))?,
 					})]),
 				}),
 				implementations: Default::default(),
@@ -163,6 +173,15 @@ impl CodegenShaderEntryPoint for PerPixelAdjust {
 			body,
 			crate_name: parsed.crate_name.clone(),
 			description: "".to_string(),
+		})?;
+
+		Ok(quote! {
+			#node_cfg
+			mod #mod_name {
+				use super::*;
+
+				#gpu_node
+			}
 		})
 	}
 }