Merge 6387f91ef9 into 877d278785

crates/bevy_core_pipeline/src/bloom/downsampling_pipeline.rs

@@ -1,8 +1,11 @@
+use core::ops::Deref;
+
 use crate::FullscreenShader;
 
 use super::{Bloom, BLOOM_TEXTURE_FORMAT};
-use bevy_asset::{load_embedded_asset, AssetServer, Handle};
+use bevy_asset::{load_embedded_asset, AssetServer};
 use bevy_ecs::{
+    error::BevyError,
     prelude::{Component, Entity},
     resource::Resource,
     system::{Commands, Query, Res, ResMut},
@@ -28,14 +31,13 @@ pub struct BloomDownsamplingPipeline {
     /// Layout with a texture, a sampler, and uniforms
     pub bind_group_layout: BindGroupLayout,
     pub sampler: Sampler,
-    /// The asset handle for the fullscreen vertex shader.
+    pub variants: SpecializedCache<RenderPipeline, BloomDownsamplingSpecializer>,
-    pub fullscreen_shader: FullscreenShader,
-    /// The fragment shader asset handle.
-    pub fragment_shader: Handle<Shader>,
 }
 
-#[derive(PartialEq, Eq, Hash, Clone)]
+pub struct BloomDownsamplingSpecializer;
-pub struct BloomDownsamplingPipelineKeys {
+
+#[derive(PartialEq, Eq, Hash, Clone, SpecializerKey)]
+pub struct BloomDownsamplingKey {
     prefilter: bool,
     first_downsample: bool,
     uniform_scale: bool,
@@ -83,27 +85,54 @@ pub fn init_bloom_downsampling_pipeline(
         ..Default::default()
     });
 
+    let fragment_shader = load_embedded_asset!(asset_server.deref(), "bloom.wgsl");
+    let base_descriptor = RenderPipelineDescriptor {
+        layout: vec![bind_group_layout.clone()],
+        vertex: fullscreen_shader.to_vertex_state(),
+        fragment: Some(FragmentState {
+            shader: fragment_shader.clone(),
+            targets: vec![Some(ColorTargetState {
+                format: BLOOM_TEXTURE_FORMAT,
+                blend: None,
+                write_mask: ColorWrites::ALL,
+            })],
+            ..default()
+        }),
+        ..default()
+    };
+
+    let variants = SpecializedCache::new(BloomDownsamplingSpecializer, None, base_descriptor);
+
     commands.insert_resource(BloomDownsamplingPipeline {
         bind_group_layout,
         sampler,
-        fullscreen_shader: fullscreen_shader.clone(),
+        variants,
-        fragment_shader: load_embedded_asset!(asset_server.as_ref(), "bloom.wgsl"),
     });
 }
 
-impl SpecializedRenderPipeline for BloomDownsamplingPipeline {
+impl Specializer<RenderPipeline> for BloomDownsamplingSpecializer {
-    type Key = BloomDownsamplingPipelineKeys;
+    type Key = BloomDownsamplingKey;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
+    fn specialize(
-        let layout = vec![self.bind_group_layout.clone()];
+        &self,
+        key: Self::Key,
+        descriptor: &mut RenderPipelineDescriptor,
+    ) -> Result<Canonical<Self::Key>, BevyError> {
+        descriptor.label = Some(if key.first_downsample {
+            "bloom_downsampling_pipeline_first".into()
+        } else {
+            "bloom_downsampling_pipeline".into()
+        });
 
-        let entry_point = if key.first_downsample {
+        let fragment = descriptor.fragment_mut()?;
+
+        fragment.entry_point = Some(if key.first_downsample {
             "downsample_first".into()
         } else {
             "downsample".into()
-        };
+        });
 
-        let mut shader_defs = vec![];
+        let shader_defs = &mut fragment.shader_defs;
 
         if key.first_downsample {
             shader_defs.push("FIRST_DOWNSAMPLE".into());
@@ -117,61 +146,36 @@ impl SpecializedRenderPipeline for BloomDownsamplingPipeline {
             shader_defs.push("UNIFORM_SCALE".into());
         }
 
-        RenderPipelineDescriptor {
+        Ok(key)
-            label: Some(
-                if key.first_downsample {
-                    "bloom_downsampling_pipeline_first"
-                } else {
-                    "bloom_downsampling_pipeline"
-                }
-                .into(),
-            ),
-            layout,
-            vertex: self.fullscreen_shader.to_vertex_state(),
-            fragment: Some(FragmentState {
-                shader: self.fragment_shader.clone(),
-                shader_defs,
-                entry_point: Some(entry_point),
-                targets: vec![Some(ColorTargetState {
-                    format: BLOOM_TEXTURE_FORMAT,
-                    blend: None,
-                    write_mask: ColorWrites::ALL,
-                })],
-            }),
-            ..default()
-        }
     }
 }
 
 pub fn prepare_downsampling_pipeline(
     mut commands: Commands,
     pipeline_cache: Res<PipelineCache>,
-    mut pipelines: ResMut<SpecializedRenderPipelines<BloomDownsamplingPipeline>>,
+    mut pipeline: ResMut<BloomDownsamplingPipeline>,
-    pipeline: Res<BloomDownsamplingPipeline>,
     views: Query<(Entity, &Bloom)>,
-) {
+) -> Result<(), BevyError> {
     for (entity, bloom) in &views {
         let prefilter = bloom.prefilter.threshold > 0.0;
 
-        let pipeline_id = pipelines.specialize(
+        let pipeline_id = pipeline.variants.specialize(
             &pipeline_cache,
-            &pipeline,
+            BloomDownsamplingKey {
-            BloomDownsamplingPipelineKeys {
                 prefilter,
                 first_downsample: false,
                 uniform_scale: bloom.scale == Vec2::ONE,
             },
-        );
+        )?;
 
-        let pipeline_first_id = pipelines.specialize(
+        let pipeline_first_id = pipeline.variants.specialize(
             &pipeline_cache,
-            &pipeline,
+            BloomDownsamplingKey {
-            BloomDownsamplingPipelineKeys {
                 prefilter,
                 first_downsample: true,
                 uniform_scale: bloom.scale == Vec2::ONE,
             },
-        );
+        )?;
 
         commands
             .entity(entity)
@@ -180,4 +184,5 @@ pub fn prepare_downsampling_pipeline(
                 main: pipeline_id,
             });
     }
+    Ok(())
 }

crates/bevy_core_pipeline/src/bloom/mod.rs

@@ -61,8 +61,6 @@ impl Plugin for BloomPlugin {
             return;
         };
         render_app
-            .init_resource::<SpecializedRenderPipelines<BloomDownsamplingPipeline>>()
-            .init_resource::<SpecializedRenderPipelines<BloomUpsamplingPipeline>>()
             .add_systems(
                 RenderStartup,
                 (

crates/bevy_core_pipeline/src/bloom/upsampling_pipeline.rs

@@ -1,10 +1,13 @@
+use core::ops::Deref;
+
 use crate::FullscreenShader;
 
 use super::{
     downsampling_pipeline::BloomUniforms, Bloom, BloomCompositeMode, BLOOM_TEXTURE_FORMAT,
 };
-use bevy_asset::{load_embedded_asset, AssetServer, Handle};
+use bevy_asset::{load_embedded_asset, AssetServer};
 use bevy_ecs::{
+    error::BevyError,
     prelude::{Component, Entity},
     resource::Resource,
     system::{Commands, Query, Res, ResMut},
@@ -28,16 +31,7 @@ pub struct UpsamplingPipelineIds {
 #[derive(Resource)]
 pub struct BloomUpsamplingPipeline {
     pub bind_group_layout: BindGroupLayout,
-    /// The asset handle for the fullscreen vertex shader.
+    pub variants: SpecializedCache<RenderPipeline, BloomUpsamplingSpecializer>,
-    pub fullscreen_shader: FullscreenShader,
-    /// The fragment shader asset handle.
-    pub fragment_shader: Handle<Shader>,
-}
-
-#[derive(PartialEq, Eq, Hash, Clone)]
-pub struct BloomUpsamplingPipelineKeys {
-    composite_mode: BloomCompositeMode,
-    final_pipeline: bool,
 }
 
 pub fn init_bloom_upscaling_pipeline(
@@ -61,17 +55,44 @@ pub fn init_bloom_upscaling_pipeline(
         ),
     );
 
+    let fragment_shader = load_embedded_asset!(asset_server.deref(), "bloom.wgsl");
+    let base_descriptor = RenderPipelineDescriptor {
+        label: Some("bloom_upsampling_pipeline".into()),
+        layout: vec![bind_group_layout.clone()],
+        vertex: fullscreen_shader.to_vertex_state(),
+        fragment: Some(FragmentState {
+            shader: fragment_shader.clone(),
+            entry_point: Some("upsample".into()),
+            ..default()
+        }),
+        ..default()
+    };
+
+    let specialized_cache =
+        SpecializedCache::new(BloomUpsamplingSpecializer, None, base_descriptor);
+
     commands.insert_resource(BloomUpsamplingPipeline {
         bind_group_layout,
-        fullscreen_shader: fullscreen_shader.clone(),
+        variants: specialized_cache,
-        fragment_shader: load_embedded_asset!(asset_server.as_ref(), "bloom.wgsl"),
     });
 }
 
-impl SpecializedRenderPipeline for BloomUpsamplingPipeline {
+pub struct BloomUpsamplingSpecializer;
-    type Key = BloomUpsamplingPipelineKeys;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
+#[derive(PartialEq, Eq, Hash, Clone, SpecializerKey)]
+pub struct BloomUpsamplingKey {
+    composite_mode: BloomCompositeMode,
+    final_pipeline: bool,
+}
+
+impl Specializer<RenderPipeline> for BloomUpsamplingSpecializer {
+    type Key = BloomUpsamplingKey;
+
+    fn specialize(
+        &self,
+        key: Self::Key,
+        descriptor: &mut RenderPipelineDescriptor,
+    ) -> Result<Canonical<Self::Key>, BevyError> {
         let texture_format = if key.final_pipeline {
             ViewTarget::TEXTURE_FORMAT_HDR
         } else {
@@ -110,61 +131,52 @@ impl SpecializedRenderPipeline for BloomUpsamplingPipeline {
             },
         };
 
-        RenderPipelineDescriptor {
+        let target = ColorTargetState {
-            label: Some("bloom_upsampling_pipeline".into()),
+            format: texture_format,
-            layout: vec![self.bind_group_layout.clone()],
+            blend: Some(BlendState {
-            vertex: self.fullscreen_shader.to_vertex_state(),
+                color: color_blend,
-            fragment: Some(FragmentState {
+                alpha: BlendComponent {
-                shader: self.fragment_shader.clone(),
+                    src_factor: BlendFactor::Zero,
-                entry_point: Some("upsample".into()),
+                    dst_factor: BlendFactor::One,
-                targets: vec![Some(ColorTargetState {
+                    operation: BlendOperation::Add,
-                    format: texture_format,
+                },
-                    blend: Some(BlendState {
-                        color: color_blend,
-                        alpha: BlendComponent {
-                            src_factor: BlendFactor::Zero,
-                            dst_factor: BlendFactor::One,
-                            operation: BlendOperation::Add,
-                        },
-                    }),
-                    write_mask: ColorWrites::ALL,
-                })],
-                ..default()
             }),
-            ..default()
+            write_mask: ColorWrites::ALL,
-        }
+        };
+
+        descriptor.fragment_mut()?.set_target(0, target);
+
+        Ok(key)
     }
 }
 
 pub fn prepare_upsampling_pipeline(
     mut commands: Commands,
     pipeline_cache: Res<PipelineCache>,
-    mut pipelines: ResMut<SpecializedRenderPipelines<BloomUpsamplingPipeline>>,
+    mut pipeline: ResMut<BloomUpsamplingPipeline>,
-    pipeline: Res<BloomUpsamplingPipeline>,
     views: Query<(Entity, &Bloom)>,
-) {
+) -> Result<(), BevyError> {
     for (entity, bloom) in &views {
-        let pipeline_id = pipelines.specialize(
+        let pipeline_id = pipeline.variants.specialize(
             &pipeline_cache,
-            &pipeline,
+            BloomUpsamplingKey {
-            BloomUpsamplingPipelineKeys {
                 composite_mode: bloom.composite_mode,
                 final_pipeline: false,
             },
-        );
+        )?;
 
-        let pipeline_final_id = pipelines.specialize(
+        let pipeline_final_id = pipeline.variants.specialize(
             &pipeline_cache,
-            &pipeline,
+            BloomUpsamplingKey {
-            BloomUpsamplingPipelineKeys {
                 composite_mode: bloom.composite_mode,
                 final_pipeline: true,
             },
-        );
+        )?;
 
         commands.entity(entity).insert(UpsamplingPipelineIds {
             id_main: pipeline_id,
             id_final: pipeline_final_id,
         });
     }
+    Ok(())
 }

crates/bevy_render/src/render_resource/pipeline.rs

@@ -7,7 +7,8 @@ use crate::{
 use alloc::borrow::Cow;
 use bevy_asset::Handle;
 use bevy_utils::WgpuWrapper;
-use core::ops::Deref;
+use core::{iter, ops::Deref};
+use thiserror::Error;
 use wgpu::{
     ColorTargetState, DepthStencilState, MultisampleState, PrimitiveState, PushConstantRange,
 };
@@ -112,6 +113,16 @@ pub struct RenderPipelineDescriptor {
     pub zero_initialize_workgroup_memory: bool,
 }
 
+#[derive(Copy, Clone, Debug, Error)]
+#[error("RenderPipelineDescriptor has no FragmentState configured")]
+pub struct NoFragmentStateError;
+
+impl RenderPipelineDescriptor {
+    pub fn fragment_mut(&mut self) -> Result<&mut FragmentState, NoFragmentStateError> {
+        self.fragment.as_mut().ok_or(NoFragmentStateError)
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Default)]
 pub struct VertexState {
     /// The compiled shader module for this stage.
@@ -137,6 +148,12 @@ pub struct FragmentState {
     pub targets: Vec<Option<ColorTargetState>>,
 }
 
+impl FragmentState {
+    pub fn set_target(&mut self, index: usize, target: ColorTargetState) {
+        filling_set_at(&mut self.targets, index, None, Some(target));
+    }
+}
+
 /// Describes a compute pipeline.
 #[derive(Clone, Debug, PartialEq, Eq, Default)]
 pub struct ComputePipelineDescriptor {
@@ -153,3 +170,11 @@ pub struct ComputePipelineDescriptor {
     /// If this is false, reading from workgroup variables before writing to them will result in garbage values.
     pub zero_initialize_workgroup_memory: bool,
 }
+
+// utility function to set a value at the specified index, extending with
+// a filler value if the index is out of bounds.
+fn filling_set_at<T: Clone>(vec: &mut Vec<T>, index: usize, filler: T, value: T) {
+    let num_to_fill = (index + 1).saturating_sub(vec.len());
+    vec.extend(iter::repeat_n(filler, num_to_fill));
+    vec[index] = value;
+}