Shader Imports. Decouple Mesh logic from PBR (#3137)

## Shader Imports

This adds "whole file" shader imports. These come in two flavors:

### Asset Path Imports

```rust
// /assets/shaders/custom.wgsl

#import "shaders/custom_material.wgsl"

[[stage(fragment)]]
fn fragment() -> [[location(0)]] vec4<f32> {
    return get_color();
}
```

```rust
// /assets/shaders/custom_material.wgsl

[[block]]
struct CustomMaterial {
    color: vec4<f32>;
};
[[group(1), binding(0)]]
var<uniform> material: CustomMaterial;
```

### Custom Path Imports

Enables defining custom import paths. These are intended to be used by crates to export shader functionality:

```rust
// bevy_pbr2/src/render/pbr.wgsl

#import bevy_pbr::mesh_view_bind_group
#import bevy_pbr::mesh_bind_group

[[block]]
struct StandardMaterial {
    base_color: vec4<f32>;
    emissive: vec4<f32>;
    perceptual_roughness: f32;
    metallic: f32;
    reflectance: f32;
    flags: u32;
};

/* rest of PBR fragment shader here */
```

```rust
impl Plugin for MeshRenderPlugin {
    fn build(&self, app: &mut bevy_app::App) {
        let mut shaders = app.world.get_resource_mut::<Assets<Shader>>().unwrap();
        shaders.set_untracked(
            MESH_BIND_GROUP_HANDLE,
            Shader::from_wgsl(include_str!("mesh_bind_group.wgsl"))
                .with_import_path("bevy_pbr::mesh_bind_group"),
        );
        shaders.set_untracked(
            MESH_VIEW_BIND_GROUP_HANDLE,
            Shader::from_wgsl(include_str!("mesh_view_bind_group.wgsl"))
                .with_import_path("bevy_pbr::mesh_view_bind_group"),
        );
```

By convention these should use rust-style module paths that start with the crate name. Ultimately we might enforce this convention.

Note that this feature implements _run time_ import resolution. Ultimately we should move the import logic into an asset preprocessor once Bevy gets support for that.

## Decouple Mesh Logic from PBR Logic via MeshRenderPlugin

This breaks out mesh rendering code from PBR material code, which improves the legibility of the code, decouples mesh logic from PBR logic, and opens the door for a future `MaterialPlugin<T: Material>` that handles all of the pipeline setup for arbitrary shader materials.

## Removed `RenderAsset<Shader>` in favor of extracting shaders into RenderPipelineCache

This simplifies the shader import implementation and removes the need to pass around `RenderAssets<Shader>`.

##  RenderCommands are now fallible

This allows us to cleanly handle pipelines+shaders not being ready yet. We can abort a render command early in these cases, preventing bevy from trying to bind group / do draw calls for pipelines that couldn't be bound. This could also be used in the future for things like "components not existing on entities yet". 

# Next Steps

* Investigate using Naga for "partial typed imports" (ex: `#import bevy_pbr::material::StandardMaterial`, which would import only the StandardMaterial struct)
* Implement `MaterialPlugin<T: Material>` for low-boilerplate custom material shaders
* Move shader import logic into the asset preprocessor once bevy gets support for that.

Fixes #3132

Author: cart

assets/shaders/custom.wgsl

@@ -0,0 +1,11 @@
+[[block]]
+struct CustomMaterial {
+    color: vec4<f32>;
+};
+[[group(1), binding(0)]]
+var<uniform> material: CustomMaterial;
+
+[[stage(fragment)]]
+fn fragment() -> [[location(0)]] vec4<f32> {
+    return material.color;
+}

assets/shaders/custom_material.wgsl

@@ -1,16 +1,6 @@
-[[block]]
-struct View {
-    view_proj: mat4x4<f32>;
-    projection: mat4x4<f32>;
-    world_position: vec3<f32>;
-};
-[[group(0), binding(0)]]
-var<uniform> view: View;
+#import bevy_pbr::mesh_view_bind_group
+#import bevy_pbr::mesh_struct
 
-[[block]]
-struct Mesh {
-    transform: mat4x4<f32>;
-};
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
@@ -26,7 +16,7 @@ struct VertexOutput {
 
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
-    let world_position = mesh.transform * vec4<f32>(vertex.position, 1.0);
+    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
 
     var out: VertexOutput;
     out.clip_position = view.view_proj * world_position;

assets/shaders/shader_defs.wgsl

@@ -46,13 +46,19 @@ impl<T: Asset> LoadedAsset<T> {
         }
     }
 
-    pub fn with_dependency(mut self, asset_path: AssetPath) -> Self {
+    pub fn add_dependency(&mut self, asset_path: AssetPath) {
         self.dependencies.push(asset_path.to_owned());
+    }
+
+    pub fn with_dependency(mut self, asset_path: AssetPath) -> Self {
+        self.add_dependency(asset_path);
         self
     }
 
-    pub fn with_dependencies(mut self, asset_paths: Vec<AssetPath<'static>>) -> Self {
-        self.dependencies.extend(asset_paths);
+    pub fn with_dependencies(mut self, mut asset_paths: Vec<AssetPath<'static>>) -> Self {
+        for asset_path in asset_paths.drain(..) {
+            self.add_dependency(asset_path);
+        }
         self
     }
 }

crates/bevy_asset/src/loader.rs

@@ -7,8 +7,8 @@ use bevy::{
     },
     math::{Vec3, Vec4},
     pbr2::{
-        DrawMesh, MeshUniform, PbrPipeline, PbrPipelineKey, SetMeshViewBindGroup,
-        SetTransformBindGroup,
+        DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
+        SetMeshViewBindGroup,
     },
     prelude::{AddAsset, App, AssetServer, Assets, GlobalTransform, Handle, Plugin, Transform},
     reflect::TypeUuid,
@@ -19,8 +19,8 @@ use bevy::{
         render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssets},
         render_component::ExtractComponentPlugin,
         render_phase::{
-            AddRenderCommand, DrawFunctions, EntityRenderCommand, RenderPhase, SetItemPipeline,
-            TrackedRenderPass,
+            AddRenderCommand, DrawFunctions, EntityRenderCommand, RenderCommandResult, RenderPhase,
+            SetItemPipeline, TrackedRenderPass,
         },
         render_resource::*,
         renderer::RenderDevice,
@@ -127,22 +127,21 @@ impl Plugin for CustomMaterialPlugin {
 }
 
 pub struct CustomPipeline {
+    mesh_pipeline: MeshPipeline,
     material_layout: BindGroupLayout,
     shader: Handle<Shader>,
-    pbr_pipeline: PbrPipeline,
 }
 
 impl SpecializedPipeline for CustomPipeline {
-    type Key = PbrPipelineKey;
+    type Key = MeshPipelineKey;
 
     fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let mut descriptor = self.pbr_pipeline.specialize(key);
-        descriptor.vertex.shader = self.shader.clone();
+        let mut descriptor = self.mesh_pipeline.specialize(key);
         descriptor.fragment.as_mut().unwrap().shader = self.shader.clone();
         descriptor.layout = Some(vec![
-            self.pbr_pipeline.view_layout.clone(),
-            self.pbr_pipeline.mesh_layout.clone(),
+            self.mesh_pipeline.view_layout.clone(),
             self.material_layout.clone(),
+            self.mesh_pipeline.mesh_layout.clone(),
         ]);
         descriptor
     }
@@ -167,8 +166,8 @@ impl FromWorld for CustomPipeline {
         });
 
         CustomPipeline {
-            pbr_pipeline: world.get_resource::<PbrPipeline>().unwrap().clone(),
-            shader: asset_server.load("shaders/custom.wgsl"),
+            mesh_pipeline: world.get_resource::<MeshPipeline>().unwrap().clone(),
+            shader: asset_server.load("shaders/custom_material.wgsl"),
             material_layout,
         }
     }
@@ -189,7 +188,7 @@ pub fn queue_custom(
         .read()
         .get_id::<DrawCustom>()
         .unwrap();
-    let key = PbrPipelineKey::from_msaa_samples(msaa.samples);
+    let key = MeshPipelineKey::from_msaa_samples(msaa.samples);
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
@@ -213,8 +212,8 @@ pub fn queue_custom(
 type DrawCustom = (
     SetItemPipeline,
     SetMeshViewBindGroup<0>,
-    SetTransformBindGroup<1>,
     SetCustomMaterialBindGroup,
+    SetMeshBindGroup<2>,
     DrawMesh,
 );
 
@@ -229,9 +228,10 @@ impl EntityRenderCommand for SetCustomMaterialBindGroup {
         item: Entity,
         (materials, query): SystemParamItem<'w, '_, Self::Param>,
         pass: &mut TrackedRenderPass<'w>,
-    ) {
+    ) -> RenderCommandResult {
         let material_handle = query.get(item).unwrap();
         let material = materials.into_inner().get(material_handle).unwrap();
-        pass.set_bind_group(2, &material.bind_group, &[]);
+        pass.set_bind_group(1, &material.bind_group, &[]);
+        RenderCommandResult::Success
     }
 }

examples/shader/custom_shader_pipelined.rs

@@ -4,8 +4,8 @@ use bevy::{
     ecs::prelude::*,
     math::Vec3,
     pbr2::{
-        DrawMesh, MeshUniform, PbrPipeline, PbrPipelineKey, SetMeshViewBindGroup,
-        SetTransformBindGroup,
+        DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
+        SetMeshViewBindGroup,
     },
     prelude::{App, AssetServer, Assets, GlobalTransform, Handle, Plugin, Transform},
     render2::{
@@ -86,39 +86,39 @@ fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
 }
 
 struct IsRedPipeline {
+    mesh_pipline: MeshPipeline,
     shader: Handle<Shader>,
-    pbr_pipeline: PbrPipeline,
 }
 
 impl FromWorld for IsRedPipeline {
     fn from_world(world: &mut World) -> Self {
         let asset_server = world.get_resource::<AssetServer>().unwrap();
-        let pbr_pipeline = world.get_resource::<PbrPipeline>().unwrap();
+        let mesh_pipeline = world.get_resource::<MeshPipeline>().unwrap();
         let shader = asset_server.load("shaders/shader_defs.wgsl");
         IsRedPipeline {
+            mesh_pipline: mesh_pipeline.clone(),
             shader,
-            pbr_pipeline: pbr_pipeline.clone(),
         }
     }
 }
 
 impl SpecializedPipeline for IsRedPipeline {
-    type Key = (IsRed, PbrPipelineKey);
+    type Key = (IsRed, MeshPipelineKey);
 
     fn specialize(&self, (is_red, pbr_pipeline_key): Self::Key) -> RenderPipelineDescriptor {
         let mut shader_defs = Vec::new();
         if is_red.0 {
             shader_defs.push("IS_RED".to_string());
         }
-        let mut descriptor = self.pbr_pipeline.specialize(pbr_pipeline_key);
+        let mut descriptor = self.mesh_pipline.specialize(pbr_pipeline_key);
         descriptor.vertex.shader = self.shader.clone();
         descriptor.vertex.shader_defs = shader_defs.clone();
         let fragment = descriptor.fragment.as_mut().unwrap();
         fragment.shader = self.shader.clone();
         fragment.shader_defs = shader_defs;
         descriptor.layout = Some(vec![
-            self.pbr_pipeline.view_layout.clone(),
-            self.pbr_pipeline.mesh_layout.clone(),
+            self.mesh_pipline.view_layout.clone(),
+            self.mesh_pipline.mesh_layout.clone(),
         ]);
         descriptor
     }
@@ -127,7 +127,7 @@ impl SpecializedPipeline for IsRedPipeline {
 type DrawIsRed = (
     SetItemPipeline,
     SetMeshViewBindGroup<0>,
-    SetTransformBindGroup<1>,
+    SetMeshBindGroup<1>,
     DrawMesh,
 );
 
@@ -144,7 +144,7 @@ fn queue_custom(
         .read()
         .get_id::<DrawIsRed>()
         .unwrap();
-    let key = PbrPipelineKey::from_msaa_samples(msaa.samples);
+    let key = MeshPipelineKey::from_msaa_samples(msaa.samples);
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);

examples/shader/shader_defs_pipelined.rs

@@ -3,7 +3,7 @@ use bevy_reflect::Reflect;
 
 // FIXME: This should probably be part of bevy_render2!
 /// Alpha mode
-#[derive(Debug, Reflect, Clone, PartialEq)]
+#[derive(Debug, Reflect, Copy, Clone, PartialEq)]
 #[reflect(Component)]
 pub enum AlphaMode {
     Opaque,

pipelined/bevy_pbr2/src/alpha.rs

@@ -16,7 +16,7 @@ use bevy_core_pipeline::{AlphaMask3d, Opaque3d, Transparent3d};
 use bevy_ecs::prelude::*;
 use bevy_reflect::TypeUuid;
 use bevy_render2::{
-    render_component::{ExtractComponentPlugin, UniformComponentPlugin},
+    render_component::ExtractComponentPlugin,
     render_graph::RenderGraph,
     render_phase::{sort_phase_system, AddRenderCommand, DrawFunctions},
     render_resource::{Shader, SpecializedPipelines},
@@ -44,14 +44,18 @@ pub struct PbrPlugin;
 impl Plugin for PbrPlugin {
     fn build(&self, app: &mut App) {
         let mut shaders = app.world.get_resource_mut::<Assets<Shader>>().unwrap();
-        let pbr_shader = Shader::from_wgsl(include_str!("render/pbr.wgsl"));
-        shaders.set_untracked(PBR_SHADER_HANDLE, pbr_shader);
-        let shadow_shader = Shader::from_wgsl(include_str!("render/depth.wgsl"));
-        shaders.set_untracked(SHADOW_SHADER_HANDLE, shadow_shader);
+        shaders.set_untracked(
+            PBR_SHADER_HANDLE,
+            Shader::from_wgsl(include_str!("render/pbr.wgsl")),
+        );
+        shaders.set_untracked(
+            SHADOW_SHADER_HANDLE,
+            Shader::from_wgsl(include_str!("render/depth.wgsl")),
+        );
 
         app.add_plugin(StandardMaterialPlugin)
+            .add_plugin(MeshRenderPlugin)
             .add_plugin(ExtractComponentPlugin::<Handle<StandardMaterial>>::default())
-            .add_plugin(UniformComponentPlugin::<MeshUniform>::default())
             .init_resource::<AmbientLight>()
             .init_resource::<DirectionalLightShadowMap>()
             .init_resource::<PointLightShadowMap>()
@@ -84,7 +88,6 @@ impl Plugin for PbrPlugin {
 
         let render_app = app.sub_app(RenderApp);
         render_app
-            .add_system_to_stage(RenderStage::Extract, render::extract_meshes)
             .add_system_to_stage(
                 RenderStage::Extract,
                 render::extract_lights.label(RenderLightSystems::ExtractLights),
@@ -97,13 +100,12 @@ impl Plugin for PbrPlugin {
                     .exclusive_system()
                     .label(RenderLightSystems::PrepareLights),
             )
-            .add_system_to_stage(RenderStage::Queue, render::queue_meshes)
             .add_system_to_stage(
                 RenderStage::Queue,
                 render::queue_shadows.label(RenderLightSystems::QueueShadows),
             )
+            .add_system_to_stage(RenderStage::Queue, queue_meshes)
             .add_system_to_stage(RenderStage::Queue, render::queue_shadow_view_bind_group)
-            .add_system_to_stage(RenderStage::Queue, render::queue_transform_bind_group)
             .add_system_to_stage(RenderStage::PhaseSort, sort_phase_system::<Shadow>)
             .init_resource::<PbrPipeline>()
             .init_resource::<ShadowPipeline>()

pipelined/bevy_pbr2/src/lib.rs

@@ -163,38 +163,44 @@ impl RenderAsset for StandardMaterial {
         material: Self::ExtractedAsset,
         (render_device, pbr_pipeline, gpu_images): &mut SystemParamItem<Self::Param>,
     ) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>> {
-        let (base_color_texture_view, base_color_sampler) = if let Some(result) =
-            pbr_pipeline.image_handle_to_texture(gpu_images, &material.base_color_texture)
+        let (base_color_texture_view, base_color_sampler) = if let Some(result) = pbr_pipeline
+            .mesh_pipeline
+            .get_image_texture(gpu_images, &material.base_color_texture)
         {
             result
         } else {
             return Err(PrepareAssetError::RetryNextUpdate(material));
         };
 
-        let (emissive_texture_view, emissive_sampler) = if let Some(result) =
-            pbr_pipeline.image_handle_to_texture(gpu_images, &material.emissive_texture)
+        let (emissive_texture_view, emissive_sampler) = if let Some(result) = pbr_pipeline
+            .mesh_pipeline
+            .get_image_texture(gpu_images, &material.emissive_texture)
         {
             result
         } else {
             return Err(PrepareAssetError::RetryNextUpdate(material));
         };
 
         let (metallic_roughness_texture_view, metallic_roughness_sampler) = if let Some(result) =
-            pbr_pipeline.image_handle_to_texture(gpu_images, &material.metallic_roughness_texture)
+            pbr_pipeline
+                .mesh_pipeline
+                .get_image_texture(gpu_images, &material.metallic_roughness_texture)
         {
             result
         } else {
             return Err(PrepareAssetError::RetryNextUpdate(material));
         };
-        let (normal_map_texture_view, normal_map_sampler) = if let Some(result) =
-            pbr_pipeline.image_handle_to_texture(gpu_images, &material.normal_map_texture)
+        let (normal_map_texture_view, normal_map_sampler) = if let Some(result) = pbr_pipeline
+            .mesh_pipeline
+            .get_image_texture(gpu_images, &material.normal_map_texture)
         {
             result
         } else {
             return Err(PrepareAssetError::RetryNextUpdate(material));
         };
-        let (occlusion_texture_view, occlusion_sampler) = if let Some(result) =
-            pbr_pipeline.image_handle_to_texture(gpu_images, &material.occlusion_texture)
+        let (occlusion_texture_view, occlusion_sampler) = if let Some(result) = pbr_pipeline
+            .mesh_pipeline
+            .get_image_texture(gpu_images, &material.occlusion_texture)
         {
             result
         } else {