Added HDR Support (pipeline specialization)

Cargo.lock

@@ -378,7 +378,7 @@ dependencies = [
 
 [[package]]
 name = "bevy_dither_post_process"
-version = "0.3.0"
+version = "0.3.1"
 dependencies = [
  "bevy",
 ]

src/components.rs

@@ -3,10 +3,13 @@ use bevy::{
     render::{
         extract_component::ExtractComponent,
         render_asset::RenderAssetUsages,
-        render_resource::{Extent3d, TextureDimension, TextureFormat, TextureUsages},
+        render_resource::{CachedRenderPipelineId, Extent3d, TextureDimension, TextureFormat, TextureUsages},
     },
 };
 
+#[derive(Component, Deref, DerefMut)]
+pub struct DitherPostProcessPipelineId(pub CachedRenderPipelineId);
+
 /// Component which, when inserted into an entity with a camera, enables the dither post-processing
 /// effect.
 #[derive(Component, ExtractComponent, Clone)]

src/lib.rs

@@ -10,9 +10,15 @@ use bevy::{
     render::{
         extract_component::ExtractComponentPlugin,
         render_graph::{RenderGraphApp, ViewNodeRunner},
-        RenderApp,
+        render_resource::{
+            PipelineCache, SpecializedRenderPipeline, SpecializedRenderPipelines, TextureFormat,
+        },
+        view::{ExtractedView, ViewTarget},
+        Render, RenderApp, RenderSet,
     },
 };
+use components::DitherPostProcessPipelineId;
+use resources::{DitherPostProcessPipeline, DitherPostProcessingPipelineKey};
 
 use crate::components::DitherPostProcessSettings;
 
@@ -38,6 +44,8 @@ impl Plugin for DitherPostProcessPlugin {
         };
 
         render_app
+            .add_systems(Render, prepare_post_processing_pipelines.in_set(RenderSet::Prepare))
+            .init_resource::<SpecializedRenderPipelines<DitherPostProcessPipeline>>()
             .add_render_graph_node::<ViewNodeRunner<nodes::DitherRenderNode>>(
                 Core3d,
                 nodes::DitherRenderLabel,
@@ -60,3 +68,27 @@ impl Plugin for DitherPostProcessPlugin {
         render_app.init_resource::<resources::DitherPostProcessPipeline>();
     }
 }
+
+fn prepare_post_processing_pipelines(
+    mut commands: Commands,
+    pipeline_cache: Res<PipelineCache>,
+    mut pipelines: ResMut<SpecializedRenderPipelines<DitherPostProcessPipeline>>,
+    post_processing_pipeline: Res<DitherPostProcessPipeline>,
+    views: Query<(Entity, &ExtractedView), With<DitherPostProcessSettings>>,
+) {
+    for (entity, view) in views.iter() {
+        let pipeline_id = pipelines.specialize(
+            &pipeline_cache,
+            &post_processing_pipeline,
+            DitherPostProcessingPipelineKey {
+                texture_format: if view.hdr {
+                    ViewTarget::TEXTURE_FORMAT_HDR
+                } else {
+                    TextureFormat::bevy_default()
+                },
+            },
+        );
+
+        commands.entity(entity).insert(DitherPostProcessPipelineId(pipeline_id));
+    }
+}

src/nodes.rs

@@ -1,14 +1,21 @@
 use bevy::{
-    ecs::query::QueryItem,
+    ecs::{query::QueryItem, system::lifetimeless::Read},
     prelude::*,
     render::{
-        render_asset::RenderAssets, render_graph::{NodeRunError, RenderGraphContext, RenderLabel, ViewNode}, render_resource::{
+        render_asset::RenderAssets,
+        render_graph::{NodeRunError, RenderGraphContext, RenderLabel, ViewNode},
+        render_resource::{
             BindGroupEntries, Operations, PipelineCache, RenderPassColorAttachment,
             RenderPassDescriptor,
-        }, renderer::RenderContext, texture::GpuImage, view::ViewTarget
+        },
+        renderer::RenderContext,
+        texture::GpuImage,
+        view::ViewTarget,
     },
 };
 
+use crate::components::DitherPostProcessPipelineId;
+
 use super::components;
 use super::resources;
 
@@ -21,20 +28,21 @@ pub struct DitherRenderNode;
 
 impl ViewNode for DitherRenderNode {
     type ViewQuery = (
-        &'static ViewTarget,
+        Read<ViewTarget>,
-        &'static components::DitherPostProcessSettings,
+        Read<components::DitherPostProcessSettings>,
+        Read<DitherPostProcessPipelineId>,
     );
 
     fn run(
         &self,
         _graph: &mut RenderGraphContext,
         render_context: &mut RenderContext,
-        (view_target, dither_post_process_settings): QueryItem<Self::ViewQuery>,
+        (view_target, dither_post_process_settings, pipeline_id): QueryItem<Self::ViewQuery>,
         world: &World,
     ) -> Result<(), NodeRunError> {
         let render_pipeline = world.resource::<resources::DitherPostProcessPipeline>();
         let pipeline_cache = world.resource::<PipelineCache>();
-        let Some(pipeline) = pipeline_cache.get_render_pipeline(render_pipeline.pipeline_id) else {
+        let Some(pipeline) = pipeline_cache.get_render_pipeline(**pipeline_id) else {
             warn!("Failed to get render pipeline from cache, skipping...");
             return Ok(());
         };

src/resources.rs

@@ -1,16 +1,10 @@
 use bevy::{
-    core_pipeline::fullscreen_vertex_shader::fullscreen_shader_vertex_state,
+    core_pipeline::fullscreen_vertex_shader, prelude::*, render::{
-    prelude::*,
-    render::{
         render_resource::{
-            binding_types::{sampler, texture_2d},
+            binding_types::{sampler, texture_2d}, BindGroupLayout, BindGroupLayoutEntries, ColorTargetState, ColorWrites, FragmentState, RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages, SpecializedRenderPipeline, TextureFormat, TextureSampleType
-            BindGroupLayout, BindGroupLayoutEntries, CachedRenderPipelineId, ColorTargetState,
-            ColorWrites, FragmentState, MultisampleState, PipelineCache, PrimitiveState,
-            RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages,
-            TextureFormat, TextureSampleType,
         },
         renderer::RenderDevice,
-    },
+    }
 };
 
 #[derive(Resource)]
@@ -18,7 +12,7 @@ pub struct DitherPostProcessPipeline {
     pub layout: BindGroupLayout,
     pub screen_sampler: Sampler,
     pub threshold_map_sampler: Sampler,
-    pub pipeline_id: CachedRenderPipelineId,
+    pub shader: Handle<Shader>,
 }
 
 impl FromWorld for DitherPostProcessPipeline {
@@ -45,35 +39,43 @@ impl FromWorld for DitherPostProcessPipeline {
             "embedded://bevy_dither_post_process/../assets/shaders/dither_post_process.wgsl",
         );
 
-        let pipeline_id =
-            world
-                .resource_mut::<PipelineCache>()
-                .queue_render_pipeline(RenderPipelineDescriptor {
-                    label: Some("dither_post_process_render_pipeline".into()),
-                    layout: vec![layout.clone()],
-                    push_constant_ranges: vec![],
-                    vertex: fullscreen_shader_vertex_state(),
-                    primitive: PrimitiveState::default(),
-                    depth_stencil: None,
-                    multisample: MultisampleState::default(),
-                    fragment: Some(FragmentState {
-                        shader,
-                        shader_defs: vec![],
-                        entry_point: "fragment".into(),
-                        targets: vec![Some(ColorTargetState {
-                            format: TextureFormat::Rgba8UnormSrgb,
-                            blend: None,
-                            write_mask: ColorWrites::ALL,
-                        })],
-                    }),
-                    zero_initialize_workgroup_memory: false,
-                });
-
         Self {
             layout,
             screen_sampler,
             threshold_map_sampler,
-            pipeline_id,
+            shader,
         }
     }
 }
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub struct DitherPostProcessingPipelineKey {
+    pub texture_format: TextureFormat,
+}
+
+impl SpecializedRenderPipeline for DitherPostProcessPipeline {
+    type Key = DitherPostProcessingPipelineKey;
+
+    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
+        RenderPipelineDescriptor {
+            label: Some("dither_post_processing".into()),
+            layout: vec![self.layout.clone()],
+            vertex: fullscreen_vertex_shader::fullscreen_shader_vertex_state(),
+            fragment: Some(FragmentState {
+                shader: self.shader.clone(),
+                shader_defs: vec![],
+                entry_point: "fragment".into(),
+                targets: vec![Some(ColorTargetState {
+                    format: key.texture_format,
+                    blend: None,
+                    write_mask: ColorWrites::ALL,
+                })],
+            }),
+            primitive: default(),
+            depth_stencil: None,
+            multisample: default(),
+            push_constant_ranges: vec![],
+            zero_initialize_workgroup_memory: false,
+        }
+    }
+}