Fix custom mesh pipelines (#3381)

# Objective

Fixes #3379 

## Solution

The custom mesh pipelines needed to be specialized on each mesh's primitive topology, as done in `queue_meshes()`

Co-authored-by: Carter Anderson <mcanders1@gmail.com>

crates/bevy_pbr/src/wireframe.rs

@@ -4,10 +4,10 @@ use bevy_app::Plugin;
 use bevy_asset::{Assets, Handle, HandleUntyped};
 use bevy_core_pipeline::Opaque3d;
 use bevy_ecs::{prelude::*, reflect::ReflectComponent};
-use bevy_reflect::Reflect;
-use bevy_reflect::TypeUuid;
+use bevy_reflect::{Reflect, TypeUuid};
 use bevy_render::{
     mesh::Mesh,
+    render_asset::RenderAssets,
     render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
     render_resource::{RenderPipelineCache, Shader, SpecializedPipeline, SpecializedPipelines},
     view::{ExtractedView, Msaa},
@@ -93,14 +93,15 @@ impl SpecializedPipeline for WireframePipeline {
 #[allow(clippy::too_many_arguments)]
 fn queue_wireframes(
     opaque_3d_draw_functions: Res<DrawFunctions<Opaque3d>>,
+    render_meshes: Res<RenderAssets<Mesh>>,
     wireframe_config: Res<WireframeConfig>,
     wireframe_pipeline: Res<WireframePipeline>,
     mut pipeline_cache: ResMut<RenderPipelineCache>,
     mut specialized_pipelines: ResMut<SpecializedPipelines<WireframePipeline>>,
     msaa: Res<Msaa>,
     mut material_meshes: QuerySet<(
-        QueryState<(Entity, &MeshUniform), With<Handle<Mesh>>>,
-        QueryState<(Entity, &MeshUniform), (With<Handle<Mesh>>, With<Wireframe>)>,
+        QueryState<(Entity, &Handle<Mesh>, &MeshUniform)>,
+        QueryState<(Entity, &Handle<Mesh>, &MeshUniform), With<Wireframe>>,
     )>,
     mut views: Query<(&ExtractedView, &mut RenderPhase<Opaque3d>)>,
 ) {
@@ -113,7 +114,11 @@ fn queue_wireframes(
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
 
-        let add_render_phase = |(entity, mesh_uniform): (Entity, &MeshUniform)| {
+        let add_render_phase =
+            |(entity, mesh_handle, mesh_uniform): (Entity, &Handle<Mesh>, &MeshUniform)| {
+                if let Some(mesh) = render_meshes.get(mesh_handle) {
+                    let key =
+                        key | MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
                     transparent_phase.add(Opaque3d {
                         entity,
                         pipeline: specialized_pipelines.specialize(
@@ -124,6 +129,7 @@ fn queue_wireframes(
                         draw_function: draw_custom,
                         distance: view_row_2.dot(mesh_uniform.transform.col(3)),
                     });
+                }
             };
 
         if wireframe_config.global {

examples/shader/shader_defs.rs

@@ -6,6 +6,7 @@ use bevy::{
     },
     prelude::*,
     render::{
+        render_asset::RenderAssets,
         render_component::{ExtractComponent, ExtractComponentPlugin},
         render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
         render_resource::{
@@ -126,13 +127,15 @@ type DrawIsRed = (
     DrawMesh,
 );
 
+#[allow(clippy::too_many_arguments)]
 fn queue_custom(
     transparent_3d_draw_functions: Res<DrawFunctions<Transparent3d>>,
+    render_meshes: Res<RenderAssets<Mesh>>,
     custom_pipeline: Res<IsRedPipeline>,
     msaa: Res<Msaa>,
     mut pipelines: ResMut<SpecializedPipelines<IsRedPipeline>>,
     mut pipeline_cache: ResMut<RenderPipelineCache>,
-    material_meshes: Query<(Entity, &MeshUniform, &IsRed), With<Handle<Mesh>>>,
+    material_meshes: Query<(Entity, &Handle<Mesh>, &MeshUniform, &IsRed)>,
     mut views: Query<(&ExtractedView, &mut RenderPhase<Transparent3d>)>,
 ) {
     let draw_custom = transparent_3d_draw_functions
@@ -143,7 +146,9 @@ fn queue_custom(
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
-        for (entity, mesh_uniform, is_red) in material_meshes.iter() {
+        for (entity, mesh_handle, mesh_uniform, is_red) in material_meshes.iter() {
+            if let Some(mesh) = render_meshes.get(mesh_handle) {
+                let key = key | MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
                 let pipeline =
                     pipelines.specialize(&mut pipeline_cache, &custom_pipeline, (*is_red, key));
                 transparent_phase.add(Transparent3d {
@@ -154,4 +159,5 @@ fn queue_custom(
                 });
             }
         }
+    }
 }

examples/shader/shader_material.rs

@@ -171,11 +171,12 @@ impl FromWorld for CustomPipeline {
 pub fn queue_custom(
     transparent_3d_draw_functions: Res<DrawFunctions<Transparent3d>>,
     materials: Res<RenderAssets<CustomMaterial>>,
+    render_meshes: Res<RenderAssets<Mesh>>,
     custom_pipeline: Res<CustomPipeline>,
     mut pipeline_cache: ResMut<RenderPipelineCache>,
     mut specialized_pipelines: ResMut<SpecializedPipelines<CustomPipeline>>,
     msaa: Res<Msaa>,
-    material_meshes: Query<(Entity, &Handle<CustomMaterial>, &MeshUniform), With<Handle<Mesh>>>,
+    material_meshes: Query<(Entity, &Handle<CustomMaterial>, &Handle<Mesh>, &MeshUniform)>,
     mut views: Query<(&ExtractedView, &mut RenderPhase<Transparent3d>)>,
 ) {
     let draw_custom = transparent_3d_draw_functions
@@ -186,8 +187,11 @@ pub fn queue_custom(
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
-        for (entity, material_handle, mesh_uniform) in material_meshes.iter() {
+        for (entity, material_handle, mesh_handle, mesh_uniform) in material_meshes.iter() {
             if materials.contains_key(material_handle) {
+                if let Some(mesh) = render_meshes.get(mesh_handle) {
+                    let key =
+                        key | MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
                     transparent_phase.add(Transparent3d {
                         entity,
                         pipeline: specialized_pipelines.specialize(
@@ -201,6 +205,7 @@ pub fn queue_custom(
                 }
             }
         }
+    }
 }
 
 type DrawCustom = (