Support AsBindGroup for 2d materials as well (#5312)

Port changes made to Material in #5053 to Material2d as well.

This is more or less an exact copy of the implementation in bevy_pbr; I
simply pretended the API existed, then copied stuff over until it
started building and the shapes example was working again.

# Objective

The changes in #5053 makes it possible to add custom materials with a lot less boiler plate. However, the implementation isn't shared with Material 2d as it's a kind of fork of the bevy_pbr version. It should be possible to use AsBindGroup on the 2d version as well.

## Solution

This makes the same kind of changes in Material2d in bevy_sprite.

This makes the following work:

```rust
//! Draws a circular purple bevy in the middle of the screen using a custom shader

use bevy::{
    prelude::*,
    reflect::TypeUuid,
    render::render_resource::{AsBindGroup, ShaderRef},
    sprite::{Material2d, Material2dPlugin, MaterialMesh2dBundle},
};

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugin(Material2dPlugin::<CustomMaterial>::default())
        .add_startup_system(setup)
        .run();
}

/// set up a simple 2D scene
fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<CustomMaterial>>,
    asset_server: Res<AssetServer>,
) {
    commands.spawn_bundle(MaterialMesh2dBundle {
        mesh: meshes.add(shape::Circle::new(50.).into()).into(),
        material: materials.add(CustomMaterial {
            color: Color::PURPLE,
            color_texture: Some(asset_server.load("branding/icon.png")),
        }),
        transform: Transform::from_translation(Vec3::new(-100., 0., 0.)),
        ..default()
    });

    commands.spawn_bundle(Camera2dBundle::default());
}

/// The Material2d trait is very configurable, but comes with sensible defaults for all methods.
/// You only need to implement functions for features that need non-default behavior. See the Material api docs for details!
impl Material2d for CustomMaterial {
    fn fragment_shader() -> ShaderRef {
        "shaders/custom_material.wgsl".into()
    }
}

// This is the struct that will be passed to your shader
#[derive(AsBindGroup, TypeUuid, Debug, Clone)]
#[uuid = "f690fdae-d598-45ab-8225-97e2a3f056e0"]
pub struct CustomMaterial {
    #[uniform(0)]
    color: Color,
    #[texture(1)]
    #[sampler(2)]
    color_texture: Option<Handle<Image>>,
}
```

Author: johanhelsing

crates/bevy_sprite/src/mesh2d/color_material.rs

@@ -1,18 +1,12 @@
 use bevy_app::{App, Plugin};
-use bevy_asset::{load_internal_asset, AssetServer, Assets, Handle, HandleUntyped};
-use bevy_ecs::system::{lifetimeless::SRes, SystemParamItem};
+use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
 use bevy_math::Vec4;
 use bevy_reflect::TypeUuid;
 use bevy_render::{
-    color::Color,
-    prelude::Shader,
-    render_asset::{PrepareAssetError, RenderAsset, RenderAssets},
-    render_resource::*,
-    renderer::RenderDevice,
-    texture::Image,
+    color::Color, prelude::Shader, render_asset::RenderAssets, render_resource::*, texture::Image,
 };
 
-use crate::{Material2d, Material2dPipeline, Material2dPlugin, MaterialMesh2dBundle};
+use crate::{Material2d, Material2dPlugin, MaterialMesh2dBundle};
 
 pub const COLOR_MATERIAL_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 3253086872234592509);
@@ -44,10 +38,13 @@ impl Plugin for ColorMaterialPlugin {
 }
 
 /// A [2d material](Material2d) that renders [2d meshes](crate::Mesh2dHandle) with a texture tinted by a uniform color
-#[derive(Debug, Clone, TypeUuid)]
+#[derive(AsBindGroup, Debug, Clone, TypeUuid)]
 #[uuid = "e228a544-e3ca-4e1e-bb9d-4d8bc1ad8c19"]
+#[uniform(0, ColorMaterialUniform)]
 pub struct ColorMaterial {
     pub color: Color,
+    #[texture(1)]
+    #[sampler(2)]
     pub texture: Option<Handle<Image>>,
 }
 
@@ -90,142 +87,28 @@ bitflags::bitflags! {
 
 /// The GPU representation of the uniform data of a [`ColorMaterial`].
 #[derive(Clone, Default, ShaderType)]
-pub struct ColorMaterialUniformData {
+pub struct ColorMaterialUniform {
     pub color: Vec4,
     pub flags: u32,
 }
 
-/// The GPU representation of a [`ColorMaterial`].
-#[derive(Debug, Clone)]
-pub struct GpuColorMaterial {
-    /// A buffer containing the [`ColorMaterialUniformData`] of the material.
-    pub buffer: Buffer,
-    /// The bind group specifying how the [`ColorMaterialUniformData`] and
-    /// the texture of the material are bound.
-    pub bind_group: BindGroup,
-    pub flags: ColorMaterialFlags,
-    pub texture: Option<Handle<Image>>,
-}
-
-impl RenderAsset for ColorMaterial {
-    type ExtractedAsset = ColorMaterial;
-    type PreparedAsset = GpuColorMaterial;
-    type Param = (
-        SRes<RenderDevice>,
-        SRes<Material2dPipeline<ColorMaterial>>,
-        SRes<RenderAssets<Image>>,
-    );
-
-    fn extract_asset(&self) -> Self::ExtractedAsset {
-        self.clone()
-    }
-
-    fn prepare_asset(
-        material: Self::ExtractedAsset,
-        (render_device, color_pipeline, gpu_images): &mut SystemParamItem<Self::Param>,
-    ) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>> {
-        let (texture_view, sampler) = if let Some(result) = color_pipeline
-            .mesh2d_pipeline
-            .get_image_texture(gpu_images, &material.texture)
-        {
-            result
-        } else {
-            return Err(PrepareAssetError::RetryNextUpdate(material));
-        };
-
+impl AsBindGroupShaderType<ColorMaterialUniform> for ColorMaterial {
+    fn as_bind_group_shader_type(&self, _images: &RenderAssets<Image>) -> ColorMaterialUniform {
         let mut flags = ColorMaterialFlags::NONE;
-        if material.texture.is_some() {
+        if self.texture.is_some() {
             flags |= ColorMaterialFlags::TEXTURE;
         }
 
-        let value = ColorMaterialUniformData {
-            color: material.color.as_linear_rgba_f32().into(),
+        ColorMaterialUniform {
+            color: self.color.as_linear_rgba_f32().into(),
             flags: flags.bits(),
-        };
-
-        let byte_buffer = [0u8; ColorMaterialUniformData::SHADER_SIZE.get() as usize];
-        let mut buffer = encase::UniformBuffer::new(byte_buffer);
-        buffer.write(&value).unwrap();
-
-        let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
-            label: Some("color_material_uniform_buffer"),
-            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
-            contents: buffer.as_ref(),
-        });
-        let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
-            entries: &[
-                BindGroupEntry {
-                    binding: 0,
-                    resource: buffer.as_entire_binding(),
-                },
-                BindGroupEntry {
-                    binding: 1,
-                    resource: BindingResource::TextureView(texture_view),
-                },
-                BindGroupEntry {
-                    binding: 2,
-                    resource: BindingResource::Sampler(sampler),
-                },
-            ],
-            label: Some("color_material_bind_group"),
-            layout: &color_pipeline.material2d_layout,
-        });
-
-        Ok(GpuColorMaterial {
-            buffer,
-            bind_group,
-            flags,
-            texture: material.texture,
-        })
+        }
     }
 }
 
 impl Material2d for ColorMaterial {
-    fn fragment_shader(_asset_server: &AssetServer) -> Option<Handle<Shader>> {
-        Some(COLOR_MATERIAL_SHADER_HANDLE.typed())
-    }
-
-    #[inline]
-    fn bind_group(render_asset: &<Self as RenderAsset>::PreparedAsset) -> &BindGroup {
-        &render_asset.bind_group
-    }
-
-    fn bind_group_layout(
-        render_device: &RenderDevice,
-    ) -> bevy_render::render_resource::BindGroupLayout {
-        render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
-            entries: &[
-                BindGroupLayoutEntry {
-                    binding: 0,
-                    visibility: ShaderStages::FRAGMENT,
-                    ty: BindingType::Buffer {
-                        ty: BufferBindingType::Uniform,
-                        has_dynamic_offset: false,
-                        min_binding_size: Some(ColorMaterialUniformData::min_size()),
-                    },
-                    count: None,
-                },
-                // Texture
-                BindGroupLayoutEntry {
-                    binding: 1,
-                    visibility: ShaderStages::FRAGMENT,
-                    ty: BindingType::Texture {
-                        multisampled: false,
-                        sample_type: TextureSampleType::Float { filterable: true },
-                        view_dimension: TextureViewDimension::D2,
-                    },
-                    count: None,
-                },
-                // Texture Sampler
-                BindGroupLayoutEntry {
-                    binding: 2,
-                    visibility: ShaderStages::FRAGMENT,
-                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
-                    count: None,
-                },
-            ],
-            label: Some("color_material_layout"),
-        })
+    fn fragment_shader() -> ShaderRef {
+        COLOR_MATERIAL_SHADER_HANDLE.typed().into()
     }
 }