From dec52a0c8f7102954243038854fa5271dc328d41 Mon Sep 17 00:00:00 2001
From: IceSentry <IceSentry@users.noreply.github.com>
Date: Tue, 8 Oct 2024 18:40:17 -0400
Subject: [PATCH] Use shader_def for oit resolve layer count (#15747)

# Objective

- Size is currently hardcoded in the shader which means it will break if
a user uses anything higher than that.

## Solution

- Use a shader_def to define the size

## Testing

Tested with the OIT example
---
 .../bevy_core_pipeline/src/oit/resolve/mod.rs | 24 +++++++++++++++----
 .../src/oit/resolve/oit_resolve.wgsl          |  3 +--
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_core_pipeline/src/oit/resolve/mod.rs b/crates/bevy_core_pipeline/src/oit/resolve/mod.rs
index 37e8a4065b9f4..ad9a8b01b7db8 100644
--- a/crates/bevy_core_pipeline/src/oit/resolve/mod.rs
+++ b/crates/bevy_core_pipeline/src/oit/resolve/mod.rs
@@ -15,7 +15,7 @@ use bevy_render::{
         BindGroup, BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries, BlendComponent,
         BlendState, CachedRenderPipelineId, ColorTargetState, ColorWrites, DownlevelFlags,
         FragmentState, MultisampleState, PipelineCache, PrimitiveState, RenderPipelineDescriptor,
-        Shader, ShaderStages, TextureFormat,
+        Shader, ShaderDefVal, ShaderStages, TextureFormat,
     },
     renderer::{RenderAdapter, RenderDevice},
     texture::BevyDefault,
@@ -121,6 +121,7 @@ pub struct OitResolvePipelineId(pub CachedRenderPipelineId);
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct OitResolvePipelineKey {
     hdr: bool,
+    layer_count: u8,
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -128,15 +129,25 @@ pub fn queue_oit_resolve_pipeline(
     mut commands: Commands,
     pipeline_cache: Res<PipelineCache>,
     resolve_pipeline: Res<OitResolvePipeline>,
-    views: Query<(Entity, &ExtractedView), With<OrderIndependentTransparencySettings>>,
+    views: Query<
+        (
+            Entity,
+            &ExtractedView,
+            &OrderIndependentTransparencySettings,
+        ),
+        With<OrderIndependentTransparencySettings>,
+    >,
     // Store the key with the id to make the clean up logic easier.
     // This also means it will always replace the entry if the key changes so nothing to clean up.
     mut cached_pipeline_id: Local<EntityHashMap<(OitResolvePipelineKey, CachedRenderPipelineId)>>,
 ) {
     let mut current_view_entities = EntityHashSet::default();
-    for (e, view) in &views {
+    for (e, view, oit_settings) in &views {
         current_view_entities.insert(e);
-        let key = OitResolvePipelineKey { hdr: view.hdr };
+        let key = OitResolvePipelineKey {
+            hdr: view.hdr,
+            layer_count: oit_settings.layer_count,
+        };
 
         if let Some((cached_key, id)) = cached_pipeline_id.get(&e) {
             if *cached_key == key {
@@ -179,7 +190,10 @@ fn specialize_oit_resolve_pipeline(
         fragment: Some(FragmentState {
             entry_point: "fragment".into(),
             shader: OIT_RESOLVE_SHADER_HANDLE,
-            shader_defs: vec![],
+            shader_defs: vec![ShaderDefVal::UInt(
+                "LAYER_COUNT".into(),
+                key.layer_count as u32,
+            )],
             targets: vec![Some(ColorTargetState {
                 format,
                 blend: Some(BlendState {
diff --git a/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl b/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
index 513000be03e35..e162793c9ec3a 100644
--- a/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
+++ b/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
@@ -12,8 +12,7 @@ struct OitFragment {
     depth: f32,
 }
 // Contains all the colors and depth for this specific fragment
-// TODO don't hardcode size
-var<private> fragment_list: array<OitFragment, 32>;
+var<private> fragment_list: array<OitFragment, #{LAYER_COUNT}>;
 
 struct FullscreenVertexOutput {
     @builtin(position) position: vec4<f32>,