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,18 +114,23 @@ 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)| {
-            transparent_phase.add(Opaque3d {
-                entity,
-                pipeline: specialized_pipelines.specialize(
-                    &mut pipeline_cache,
-                    &wireframe_pipeline,
-                    key,
-                ),
-                draw_function: draw_custom,
-                distance: view_row_2.dot(mesh_uniform.transform.col(3)),
-            });
-        };
+        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(
+                            &mut pipeline_cache,
+                            &wireframe_pipeline,
+                            key,
+                        ),
+                        draw_function: draw_custom,
+                        distance: view_row_2.dot(mesh_uniform.transform.col(3)),
+                    });
+                }
+            };
 
         if wireframe_config.global {
             material_meshes.q0().iter().for_each(add_render_phase);

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,15 +146,18 @@ 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() {
-            let pipeline =
-                pipelines.specialize(&mut pipeline_cache, &custom_pipeline, (*is_red, key));
-            transparent_phase.add(Transparent3d {
-                entity,
-                pipeline,
-                draw_function: draw_custom,
-                distance: view_row_2.dot(mesh_uniform.transform.col(3)),
-            });
+        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 {
+                    entity,
+                    pipeline,
+                    draw_function: draw_custom,
+                    distance: view_row_2.dot(mesh_uniform.transform.col(3)),
+                });
+            }
         }
     }
 }

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,18 +187,22 @@ 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) {
-                transparent_phase.add(Transparent3d {
-                    entity,
-                    pipeline: specialized_pipelines.specialize(
-                        &mut pipeline_cache,
-                        &custom_pipeline,
-                        key,
-                    ),
-                    draw_function: draw_custom,
-                    distance: view_row_2.dot(mesh_uniform.transform.col(3)),
-                });
+                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(
+                            &mut pipeline_cache,
+                            &custom_pipeline,
+                            key,
+                        ),
+                        draw_function: draw_custom,
+                        distance: view_row_2.dot(mesh_uniform.transform.col(3)),
+                    });
+                }
             }
         }
     }