From 64faadb9325938d255386c8ee535d0ee03e88c82 Mon Sep 17 00:00:00 2001
From: IceSentry <IceSentry@users.noreply.github.com>
Date: Fri, 3 Nov 2023 19:38:50 -0400
Subject: [PATCH] Fix gizmo crash when prepass enabled (#10360)

# Objective

- Fix gizmo crash when prepass enabled

## Solution

- Add the prepass to the view key

Fixes: https://github.com/bevyengine/bevy/issues/10347
---
 crates/bevy_gizmos/src/pipeline_3d.rs | 48 ++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_gizmos/src/pipeline_3d.rs b/crates/bevy_gizmos/src/pipeline_3d.rs
index 416d7e9a9992f..2cb016753513b 100644
--- a/crates/bevy_gizmos/src/pipeline_3d.rs
+++ b/crates/bevy_gizmos/src/pipeline_3d.rs
@@ -4,10 +4,14 @@ use crate::{
 };
 use bevy_app::{App, Plugin};
 use bevy_asset::Handle;
-use bevy_core_pipeline::core_3d::{Transparent3d, CORE_3D_DEPTH_FORMAT};
+use bevy_core_pipeline::{
+    core_3d::{Transparent3d, CORE_3D_DEPTH_FORMAT},
+    prepass::{DeferredPrepass, DepthPrepass, MotionVectorPrepass, NormalPrepass},
+};
 
 use bevy_ecs::{
     prelude::Entity,
+    query::Has,
     schedule::IntoSystemConfigs,
     system::{Query, Res, ResMut, Resource},
     world::{FromWorld, World},
@@ -69,7 +73,7 @@ impl FromWorld for LineGizmoPipeline {
 
 #[derive(PartialEq, Eq, Hash, Clone)]
 struct LineGizmoPipelineKey {
-    mesh_key: MeshPipelineKey,
+    view_key: MeshPipelineKey,
     strip: bool,
     perspective: bool,
 }
@@ -87,7 +91,7 @@ impl SpecializedRenderPipeline for LineGizmoPipeline {
             shader_defs.push("PERSPECTIVE".into());
         }
 
-        let format = if key.mesh_key.contains(MeshPipelineKey::HDR) {
+        let format = if key.view_key.contains(MeshPipelineKey::HDR) {
             ViewTarget::TEXTURE_FORMAT_HDR
         } else {
             TextureFormat::bevy_default()
@@ -95,7 +99,7 @@ impl SpecializedRenderPipeline for LineGizmoPipeline {
 
         let view_layout = self
             .mesh_pipeline
-            .get_view_layout(key.mesh_key.into())
+            .get_view_layout(key.view_key.into())
             .clone();
 
         let layout = vec![view_layout, self.uniform_layout.clone()];
@@ -127,7 +131,7 @@ impl SpecializedRenderPipeline for LineGizmoPipeline {
                 bias: DepthBiasState::default(),
             }),
             multisample: MultisampleState {
-                count: key.mesh_key.msaa_samples(),
+                count: key.view_key.msaa_samples(),
                 mask: !0,
                 alpha_to_coverage_enabled: false,
             },
@@ -158,19 +162,47 @@ fn queue_line_gizmos_3d(
         &ExtractedView,
         &mut RenderPhase<Transparent3d>,
         Option<&RenderLayers>,
+        (
+            Has<NormalPrepass>,
+            Has<DepthPrepass>,
+            Has<MotionVectorPrepass>,
+            Has<DeferredPrepass>,
+        ),
     )>,
 ) {
     let draw_function = draw_functions.read().get_id::<DrawLineGizmo3d>().unwrap();
 
-    for (view, mut transparent_phase, render_layers) in &mut views {
+    for (
+        view,
+        mut transparent_phase,
+        render_layers,
+        (normal_prepass, depth_prepass, motion_vector_prepass, deferred_prepass),
+    ) in &mut views
+    {
         let render_layers = render_layers.copied().unwrap_or_default();
         if !config.render_layers.intersects(&render_layers) {
             continue;
         }
 
-        let mesh_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
+        let mut view_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
             | MeshPipelineKey::from_hdr(view.hdr);
 
+        if normal_prepass {
+            view_key |= MeshPipelineKey::NORMAL_PREPASS;
+        }
+
+        if depth_prepass {
+            view_key |= MeshPipelineKey::DEPTH_PREPASS;
+        }
+
+        if motion_vector_prepass {
+            view_key |= MeshPipelineKey::MOTION_VECTOR_PREPASS;
+        }
+
+        if deferred_prepass {
+            view_key |= MeshPipelineKey::DEFERRED_PREPASS;
+        }
+
         for (entity, handle) in &line_gizmos {
             let Some(line_gizmo) = line_gizmo_assets.get(handle) else {
                 continue;
@@ -180,7 +212,7 @@ fn queue_line_gizmos_3d(
                 &pipeline_cache,
                 &pipeline,
                 LineGizmoPipelineKey {
-                    mesh_key,
+                    view_key,
                     strip: line_gizmo.strip,
                     perspective: config.line_perspective,
                 },