diff --git a/crates/bevy_anti_alias/src/taa/mod.rs b/crates/bevy_anti_alias/src/taa/mod.rs index e396e9accd67b..377de519c1560 100644 --- a/crates/bevy_anti_alias/src/taa/mod.rs +++ b/crates/bevy_anti_alias/src/taa/mod.rs @@ -1,5 +1,5 @@ use bevy_app::{App, Plugin}; -use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer, Handle}; +use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer}; use bevy_camera::{Camera, Camera3d}; use bevy_core_pipeline::{ core_3d::graph::{Core3d, Node3d}, @@ -8,6 +8,7 @@ use bevy_core_pipeline::{ }; use bevy_diagnostic::FrameCount; use bevy_ecs::{ + error::BevyError, prelude::{Component, Entity, ReflectComponent}, query::{QueryItem, With}, resource::Resource, @@ -25,11 +26,11 @@ use bevy_render::{ render_resource::{ binding_types::{sampler, texture_2d, texture_depth_2d}, BindGroupEntries, BindGroupLayoutDescriptor, BindGroupLayoutEntries, - CachedRenderPipelineId, ColorTargetState, ColorWrites, FilterMode, FragmentState, - Operations, PipelineCache, RenderPassColorAttachment, RenderPassDescriptor, - RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages, - SpecializedRenderPipeline, SpecializedRenderPipelines, TextureDescriptor, TextureDimension, - TextureFormat, TextureSampleType, TextureUsages, + CachedRenderPipelineId, Canonical, ColorTargetState, ColorWrites, FilterMode, + FragmentState, Operations, PipelineCache, RenderPassColorAttachment, RenderPassDescriptor, + RenderPipeline, RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, + ShaderStages, Specializer, SpecializerKey, TextureDescriptor, TextureDimension, + TextureFormat, TextureSampleType, TextureUsages, Variants, }, renderer::{RenderContext, RenderDevice}, sync_component::SyncComponentPlugin, @@ -38,7 +39,6 @@ use bevy_render::{ view::{ExtractedView, Msaa, ViewTarget}, ExtractSchedule, MainWorld, Render, RenderApp, RenderStartup, RenderSystems, }; -use bevy_shader::Shader; use bevy_utils::default; use tracing::warn; @@ -58,7 +58,6 @@ impl Plugin for TemporalAntiAliasPlugin { return; }; render_app - .init_resource::>() .add_systems(RenderStartup, init_taa_pipeline) .add_systems(ExtractSchedule, extract_taa_settings) .add_systems( @@ -239,8 +238,7 @@ struct TaaPipeline { taa_bind_group_layout: BindGroupLayoutDescriptor, nearest_sampler: Sampler, linear_sampler: Sampler, - fullscreen_shader: FullscreenShader, - fragment_shader: Handle, + variants: Variants, } fn init_taa_pipeline( @@ -283,61 +281,68 @@ fn init_taa_pipeline( ), ); + let fragment_shader = load_embedded_asset!(asset_server.as_ref(), "taa.wgsl"); + + let variants = Variants::new( + TaaPipelineSpecializer, + RenderPipelineDescriptor { + label: Some("taa_pipeline".into()), + layout: vec![taa_bind_group_layout.clone()], + vertex: fullscreen_shader.to_vertex_state(), + fragment: Some(FragmentState { + shader: fragment_shader, + ..default() + }), + ..default() + }, + ); + commands.insert_resource(TaaPipeline { taa_bind_group_layout, nearest_sampler, linear_sampler, - fullscreen_shader: fullscreen_shader.clone(), - fragment_shader: load_embedded_asset!(asset_server.as_ref(), "taa.wgsl"), + variants, }); } -#[derive(PartialEq, Eq, Hash, Clone)] +struct TaaPipelineSpecializer; + +#[derive(PartialEq, Eq, Hash, Clone, SpecializerKey)] struct TaaPipelineKey { hdr: bool, reset: bool, } -impl SpecializedRenderPipeline for TaaPipeline { +impl Specializer for TaaPipelineSpecializer { type Key = TaaPipelineKey; - fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { - let mut shader_defs = vec![]; - + fn specialize( + &self, + key: Self::Key, + descriptor: &mut RenderPipelineDescriptor, + ) -> Result, BevyError> { + let fragment = descriptor.fragment_mut()?; let format = if key.hdr { - shader_defs.push("TONEMAP".into()); + fragment.shader_defs.push("TONEMAP".into()); ViewTarget::TEXTURE_FORMAT_HDR } else { TextureFormat::bevy_default() }; if key.reset { - shader_defs.push("RESET".into()); + fragment.shader_defs.push("RESET".into()); } - RenderPipelineDescriptor { - label: Some("taa_pipeline".into()), - layout: vec![self.taa_bind_group_layout.clone()], - vertex: self.fullscreen_shader.to_vertex_state(), - fragment: Some(FragmentState { - shader: self.fragment_shader.clone(), - shader_defs, - targets: vec![ - Some(ColorTargetState { - format, - blend: None, - write_mask: ColorWrites::ALL, - }), - Some(ColorTargetState { - format, - blend: None, - write_mask: ColorWrites::ALL, - }), - ], - ..default() - }), - ..default() - } + let color_target_state = ColorTargetState { + format, + blend: None, + write_mask: ColorWrites::ALL, + }; + + fragment.set_target(0, color_target_state.clone()); + fragment.set_target(1, color_target_state); + + Ok(key) } } @@ -455,25 +460,30 @@ pub struct TemporalAntiAliasPipelineId(CachedRenderPipelineId); fn prepare_taa_pipelines( mut commands: Commands, pipeline_cache: Res, - mut pipelines: ResMut>, - pipeline: Res, + mut pipeline: ResMut, views: Query<(Entity, &ExtractedView, &TemporalAntiAliasing)>, -) { +) -> Result<(), BevyError> { for (entity, view, taa_settings) in &views { let mut pipeline_key = TaaPipelineKey { hdr: view.hdr, reset: taa_settings.reset, }; - let pipeline_id = pipelines.specialize(&pipeline_cache, &pipeline, pipeline_key.clone()); + let pipeline_id = pipeline + .variants + .specialize(&pipeline_cache, pipeline_key.clone())?; // Prepare non-reset pipeline anyways - it will be necessary next frame if pipeline_key.reset { pipeline_key.reset = false; - pipelines.specialize(&pipeline_cache, &pipeline, pipeline_key); + pipeline + .variants + .specialize(&pipeline_cache, pipeline_key)?; } commands .entity(entity) .insert(TemporalAntiAliasPipelineId(pipeline_id)); } + + Ok(()) }