349 lines
11 KiB
Rust
349 lines
11 KiB
Rust
use crate::crate_ident::CrateIdent;
|
|
use crate::parsing::{Input, NodeFnAttributes, ParsedField, ParsedFieldType, ParsedNodeFn, ParsedValueSource, RegularParsedField};
|
|
use crate::shader_nodes::{SHADER_NODES_FEATURE_GATE, ShaderCodegen, ShaderNodeType, ShaderTokens};
|
|
use convert_case::{Case, Casing};
|
|
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::{LitStr, PatIdent, Type, parse_quote};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct PerPixelAdjust {}
|
|
|
|
impl Parse for PerPixelAdjust {
|
|
fn parse(_input: ParseStream) -> syn::Result<Self> {
|
|
Ok(Self {})
|
|
}
|
|
}
|
|
|
|
impl ShaderCodegen for PerPixelAdjust {
|
|
fn codegen(&self, crate_ident: &CrateIdent, parsed: &ParsedNodeFn) -> syn::Result<ShaderTokens> {
|
|
let fn_name = &parsed.fn_name;
|
|
|
|
let mut params;
|
|
let has_uniform;
|
|
{
|
|
// categorize params
|
|
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: ty.to_token_stream(),
|
|
param_type: ParamType::Uniform,
|
|
}),
|
|
ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. }) => {
|
|
let param = Param {
|
|
ident: Cow::Owned(format_ident!("image_{}", &ident.ident)),
|
|
ty: quote!(Image2d),
|
|
param_type: ParamType::Image { binding: 0 },
|
|
};
|
|
Ok(param)
|
|
}
|
|
}
|
|
})
|
|
.collect::<syn::Result<Vec<_>>>()?;
|
|
|
|
has_uniform = params.iter().any(|p| matches!(p.param_type, ParamType::Uniform));
|
|
|
|
// assign image bindings
|
|
// if an arg_buffer exists, bindings for images start at 1 to leave 0 for arg buffer
|
|
let mut binding_cnt = if has_uniform { 1 } else { 0 };
|
|
for p in params.iter_mut() {
|
|
match &mut p.param_type {
|
|
ParamType::Image { binding } => {
|
|
*binding = binding_cnt;
|
|
binding_cnt += 1;
|
|
}
|
|
ParamType::Uniform => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
let entry_point_mod = format_ident!("{}_gpu", fn_name);
|
|
let entry_point_name_ident = format_ident!("ENTRY_POINT_NAME");
|
|
let entry_point_name = quote!(#entry_point_mod::#entry_point_name_ident);
|
|
let uniform_struct_ident = format_ident!("Uniform");
|
|
let uniform_struct = quote!(#entry_point_mod::#uniform_struct_ident);
|
|
let shader_node_mod = format_ident!("{}_shader_node", fn_name);
|
|
|
|
let codegen = PerPixelAdjustCodegen {
|
|
crate_ident,
|
|
parsed,
|
|
params,
|
|
has_uniform,
|
|
entry_point_mod,
|
|
entry_point_name_ident,
|
|
entry_point_name,
|
|
uniform_struct_ident,
|
|
uniform_struct,
|
|
shader_node_mod,
|
|
};
|
|
|
|
Ok(ShaderTokens {
|
|
shader_entry_point: codegen.codegen_shader_entry_point()?,
|
|
gpu_node: codegen.codegen_gpu_node()?,
|
|
})
|
|
}
|
|
}
|
|
|
|
pub struct PerPixelAdjustCodegen<'a> {
|
|
crate_ident: &'a CrateIdent,
|
|
parsed: &'a ParsedNodeFn,
|
|
params: Vec<Param<'a>>,
|
|
has_uniform: bool,
|
|
entry_point_mod: Ident,
|
|
entry_point_name_ident: Ident,
|
|
entry_point_name: TokenStream,
|
|
uniform_struct_ident: Ident,
|
|
uniform_struct: TokenStream,
|
|
shader_node_mod: Ident,
|
|
}
|
|
|
|
impl PerPixelAdjustCodegen<'_> {
|
|
fn codegen_shader_entry_point(&self) -> syn::Result<TokenStream> {
|
|
let fn_name = &self.parsed.fn_name;
|
|
let gcore_shaders = self.crate_ident.gcore_shaders()?;
|
|
let reexport = quote!(#gcore_shaders::shaders::__private);
|
|
|
|
let uniform_members = self
|
|
.params
|
|
.iter()
|
|
.filter_map(|Param { ident, ty, param_type }| match param_type {
|
|
ParamType::Image { .. } => None,
|
|
ParamType::Uniform => Some(quote! {#ident: #ty}),
|
|
})
|
|
.collect::<Vec<_>>();
|
|
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()
|
|
.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 = self
|
|
.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!(());
|
|
|
|
let entry_point_mod = &self.entry_point_mod;
|
|
let entry_point_name = &self.entry_point_name_ident;
|
|
Ok(quote! {
|
|
pub mod #entry_point_mod {
|
|
use super::*;
|
|
use #gcore_shaders::color::Color;
|
|
use #reexport::glam::{Vec4, Vec4Swizzles};
|
|
use #reexport::spirv_std::spirv;
|
|
use #reexport::spirv_std::image::{Image2d, ImageWithMethods};
|
|
use #reexport::spirv_std::image::sample_with::lod;
|
|
|
|
pub const #entry_point_name: &str = core::concat!(core::module_path!(), "::entry_point");
|
|
|
|
#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: &UniformBuffer,
|
|
#(#image_params),*
|
|
) {
|
|
let uniform = <Uniform as #gcore_shaders::shaders::buffer_struct::BufferStruct>::read(*uniform);
|
|
let texel_coord = frag_coord.xy().as_uvec2();
|
|
let color: Color = #fn_name(#context, #(#call_args),*);
|
|
*color_out = color.to_vec4();
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
fn codegen_gpu_node(&self) -> syn::Result<TokenStream> {
|
|
let gcore = self.crate_ident.gcore()?;
|
|
let raster_types = self.crate_ident.raster_types()?;
|
|
let wgpu_executor = self.crate_ident.wgpu_executor()?;
|
|
|
|
// adapt fields for gpu node
|
|
let raster_gpu: Type = parse_quote!(#gcore::table::Table<#raster_types::Raster<#raster_types::GPU>>);
|
|
let mut fields = self
|
|
.parsed
|
|
.fields
|
|
.iter()
|
|
.map(|f| match &f.ty {
|
|
ParsedFieldType::Regular(reg @ RegularParsedField { gpu_image: true, .. }) => Ok(ParsedField {
|
|
pat_ident: PatIdent {
|
|
mutability: None,
|
|
by_ref: None,
|
|
..f.pat_ident.clone()
|
|
},
|
|
ty: ParsedFieldType::Regular(RegularParsedField {
|
|
ty: raster_gpu.clone(),
|
|
implementations: Punctuated::default(),
|
|
..reg.clone()
|
|
}),
|
|
..f.clone()
|
|
}),
|
|
ParsedFieldType::Regular(RegularParsedField { gpu_image: false, .. }) => Ok(ParsedField {
|
|
pat_ident: PatIdent {
|
|
mutability: None,
|
|
by_ref: None,
|
|
..f.pat_ident.clone()
|
|
},
|
|
..f.clone()
|
|
}),
|
|
ParsedFieldType::Node { .. } => Err(syn::Error::new_spanned(&f.pat_ident, "PerPixelAdjust shader nodes cannot accept other nodes as generics")),
|
|
})
|
|
.collect::<syn::Result<Vec<_>>>()?;
|
|
|
|
// insert wgpu_executor field
|
|
let executor = format_ident!("__wgpu_executor");
|
|
fields.push(ParsedField {
|
|
pat_ident: PatIdent {
|
|
attrs: vec![],
|
|
by_ref: None,
|
|
mutability: None,
|
|
ident: parse_quote!(#executor),
|
|
subpat: None,
|
|
},
|
|
name: None,
|
|
description: "".to_string(),
|
|
widget_override: Default::default(),
|
|
ty: ParsedFieldType::Regular(RegularParsedField {
|
|
ty: parse_quote!(&'a WgpuExecutor),
|
|
exposed: true,
|
|
value_source: ParsedValueSource::Scope(LitStr::new("wgpu-executor", Span::call_site())),
|
|
number_soft_min: None,
|
|
number_soft_max: None,
|
|
number_hard_min: None,
|
|
number_hard_max: None,
|
|
number_mode_range: None,
|
|
implementations: Default::default(),
|
|
gpu_image: false,
|
|
}),
|
|
number_display_decimal_places: None,
|
|
number_step: None,
|
|
unit: None,
|
|
is_data_field: false,
|
|
});
|
|
|
|
// find exactly one gpu_image field, runtime doesn't support more than 1 atm
|
|
let gpu_image_field = {
|
|
let mut iter = fields.iter().filter(|f| matches!(f.ty, ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. })));
|
|
match (iter.next(), iter.next()) {
|
|
(Some(v), None) => Ok(v),
|
|
(Some(_), Some(more)) => Err(syn::Error::new_spanned(&more.pat_ident, "No more than one parameter must be annotated with `#[gpu_image]`")),
|
|
(None, _) => Err(syn::Error::new_spanned(&self.parsed.fn_name, "At least one parameter must be annotated with `#[gpu_image]`")),
|
|
}?
|
|
};
|
|
let gpu_image = &gpu_image_field.pat_ident.ident;
|
|
|
|
// uniform buffer struct construction
|
|
let has_uniform = self.has_uniform;
|
|
let uniform_buffer = if has_uniform {
|
|
let uniform_struct = &self.uniform_struct;
|
|
let uniform_members = self
|
|
.params
|
|
.iter()
|
|
.filter_map(|p| match p.param_type {
|
|
ParamType::Image { .. } => None,
|
|
ParamType::Uniform => Some(p.ident.as_ref()),
|
|
})
|
|
.collect::<Vec<_>>();
|
|
quote!(Some(&super::#uniform_struct {
|
|
#(#uniform_members),*
|
|
}))
|
|
} else {
|
|
// explicit generics placed here cause it's easier than explicitly writing `run_per_pixel_adjust::<()>`
|
|
quote!(Option::<&()>::None)
|
|
};
|
|
|
|
// node function body
|
|
let entry_point_name = &self.entry_point_name;
|
|
let body = quote! {
|
|
{
|
|
#executor.shader_runtime.run_per_pixel_adjust(&::wgpu_executor::shader_runtime::per_pixel_adjust_runtime::Shaders {
|
|
wgsl_shader: crate::WGSL_SHADER,
|
|
fragment_shader_name: super::#entry_point_name,
|
|
has_uniform: #has_uniform,
|
|
}, #gpu_image, #uniform_buffer).await
|
|
}
|
|
};
|
|
|
|
// call node codegen
|
|
let display_name = self.parsed.attributes.display_name.clone();
|
|
let display_name = display_name.unwrap_or_else(|| LitStr::new(&self.shader_node_mod.to_string().strip_suffix("_shader_node").unwrap().to_case(Case::Title), Span::call_site()));
|
|
let display_name = LitStr::new(&format!("{} GPU", display_name.value()), display_name.span());
|
|
let mut parsed_node_fn = ParsedNodeFn {
|
|
vis: self.parsed.vis.clone(),
|
|
attributes: NodeFnAttributes {
|
|
display_name: Some(display_name),
|
|
shader_node: Some(ShaderNodeType::ShaderNode),
|
|
..self.parsed.attributes.clone()
|
|
},
|
|
fn_name: self.shader_node_mod.clone(),
|
|
struct_name: format_ident!("{}", self.shader_node_mod.to_string().to_case(Case::Pascal)),
|
|
mod_name: self.shader_node_mod.clone(),
|
|
fn_generics: vec![parse_quote!('a: 'n)],
|
|
where_clause: None,
|
|
input: Input {
|
|
pat_ident: self.parsed.input.pat_ident.clone(),
|
|
ty: parse_quote!(impl #gcore::context::Ctx),
|
|
implementations: Default::default(),
|
|
context_features: self.parsed.input.context_features.clone(),
|
|
},
|
|
output_type: raster_gpu,
|
|
is_async: true,
|
|
fields,
|
|
body,
|
|
description: self.parsed.description.clone(),
|
|
};
|
|
parsed_node_fn.replace_impl_trait_in_input();
|
|
let gpu_node_impl = crate::codegen::generate_node_code(self.crate_ident, &parsed_node_fn)?;
|
|
|
|
// wrap node in `mod #gpu_node_mod`
|
|
let shader_node_mod = &self.shader_node_mod;
|
|
Ok(quote! {
|
|
#[cfg(feature = #SHADER_NODES_FEATURE_GATE)]
|
|
mod #shader_node_mod {
|
|
use super::*;
|
|
use #wgpu_executor::WgpuExecutor;
|
|
|
|
#gpu_node_impl
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
struct Param<'a> {
|
|
ident: Cow<'a, Ident>,
|
|
ty: TokenStream,
|
|
param_type: ParamType,
|
|
}
|
|
|
|
enum ParamType {
|
|
Image { binding: u32 },
|
|
Uniform,
|
|
}
|