Broaden RenderAsset API

benches/benches/bevy_render/extract_render_asset.rs

@@ -1,10 +1,10 @@
 use bevy_app::{App, AppLabel};
-use bevy_asset::{Asset, AssetApp, AssetEvent, AssetId, Assets, RenderAssetUsages};
+use bevy_asset::{Asset, AssetApp, AssetEvent, AssetId, Assets};
 use bevy_ecs::prelude::*;
 use bevy_reflect::TypePath;
 use bevy_render::{
     extract_plugin::ExtractPlugin,
-    render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin},
+    render_asset::{AlreadyTaken, PrepareAssetError, RenderAsset, RenderAssetPlugin},
     RenderApp,
 };
 use criterion::{criterion_group, BenchmarkId, Criterion, Throughput};
@@ -19,8 +19,13 @@ impl RenderAsset for DummyRenderAsset {
     type SourceAsset = DummyAsset;
     type Param = ();
 
-    fn asset_usage(_: &Self::SourceAsset) -> RenderAssetUsages {
-        RenderAssetUsages::RENDER_WORLD
+    type Extracted = Self::SourceAsset;
+
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
+        _previous_gpu_asset: Option<&Self>,
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        Some(Ok(source_asset.clone()))
     }
 
     fn prepare_asset(

crates/bevy_asset/src/render_asset.rs

@@ -50,3 +50,28 @@ impl Default for RenderAssetUsages {
         RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD
     }
 }
+
+impl RenderAssetUsages {
+    /// Function that conditionally extracts data from an asset based on the render asset usage flags.
+    ///
+    /// If the asset is not intended for the render world, `None` is returned.
+    ///
+    /// If it's only intended for the render world `take` is used. Otherwise `clone` is used.
+    pub fn extract<I, O>(
+        self,
+        input: &mut I,
+        take: impl FnOnce(&mut I) -> O,
+        clone: impl FnOnce(&mut I) -> O,
+    ) -> Option<O> {
+        if self.contains(RenderAssetUsages::RENDER_WORLD) {
+            let res = if self == Self::RENDER_WORLD {
+                take(input)
+            } else {
+                clone(input)
+            };
+            Some(res)
+        } else {
+            None
+        }
+    }
+}

crates/bevy_gizmos_render/src/lib.rs

@@ -70,6 +70,7 @@ use {
 
 use bevy_render::{
     extract_resource::{ExtractResource, ExtractResourcePlugin},
+    render_asset::AlreadyTaken,
     render_resource::{BindGroupLayoutDescriptor, PipelineCache, VertexAttribute, VertexStepMode},
 };
 
@@ -255,12 +256,21 @@ impl RenderAsset for GpuLineGizmo {
     type SourceAsset = GizmoAsset;
     type Param = SRes<RenderDevice>;
 
+    type Extracted = Self::SourceAsset;
+
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
+        _previous_gpu_asset: Option<&Self>,
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        Some(Ok(source_asset.clone()))
+    }
+
     fn prepare_asset(
-        gizmo: Self::SourceAsset,
+        gizmo: Self::Extracted,
         _: AssetId<Self::SourceAsset>,
         render_device: &mut SystemParamItem<Self::Param>,
         _: Option<&Self>,
-    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
+    ) -> Result<Self, PrepareAssetError<Self::Extracted>> {
         let list_position_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
             usage: BufferUsages::VERTEX,
             label: Some("LineGizmo Position Buffer"),

crates/bevy_pbr/src/medium.rs

@@ -9,7 +9,7 @@ use bevy_ecs::{
 };
 use bevy_math::{ops, Vec4};
 use bevy_render::{
-    render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin},
+    render_asset::{AlreadyTaken, PrepareAssetError, RenderAsset, RenderAssetPlugin},
     render_resource::{
         Extent3d, FilterMode, Sampler, SamplerDescriptor, Texture, TextureDataOrder,
         TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, TextureView,
@@ -66,12 +66,21 @@ impl RenderAsset for GpuScatteringMedium {
 
     type Param = (Res<'static, RenderDevice>, Res<'static, RenderQueue>);
 
+    type Extracted = Self::SourceAsset;
+
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
+        _previous_gpu_asset: Option<&Self>,
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        Some(Ok(source_asset.clone()))
+    }
+
     fn prepare_asset(
-        source_asset: Self::SourceAsset,
+        source_asset: Self::Extracted,
         _asset_id: AssetId<Self::SourceAsset>,
         (render_device, render_queue): &mut SystemParamItem<Self::Param>,
         _previous_asset: Option<&Self>,
-    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
+    ) -> Result<Self, PrepareAssetError<Self::Extracted>> {
         let mut density: Vec<Vec4> =
             Vec::with_capacity(2 * source_asset.falloff_resolution as usize);
 

crates/bevy_pbr/src/render/mesh.rs

@@ -44,7 +44,7 @@ use bevy_render::batching::gpu_preprocessing::PreviousInstanceInputUniformBuffer
 use bevy_render::impl_atomic_pod;
 use bevy_render::mesh::allocator::{MeshSlabId, MeshSlabs};
 use bevy_render::mesh::morph::{
-    MorphTargetImage, MorphTargetsResource, RenderMorphTargetAllocator,
+    MorphTargetImages, MorphTargetsResource, RenderMorphTargetAllocator,
 };
 use bevy_render::{
     batching::{
@@ -3969,7 +3969,9 @@ fn prepare_mesh_bind_groups_for_phase(
     if weights_uniform.current_buffer.buffer().is_some() {
         match (render_morph_target_allocator, &mut groups.morph_targets) {
             (
-                RenderMorphTargetAllocator::Image { mesh_id_to_image },
+                RenderMorphTargetAllocator::Image {
+                    morph_target_images,
+                },
                 &mut MeshMorphTargetBindGroups::Uniform(ref mut morph_targets),
             ) => {
                 prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
@@ -3980,7 +3982,7 @@ fn prepare_mesh_bind_groups_for_phase(
                     pipeline_cache,
                     skins_uniform,
                     weights_uniform,
-                    mesh_id_to_image,
+                    morph_target_images,
                     morph_targets,
                 );
             }
@@ -4039,7 +4041,7 @@ fn prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
     pipeline_cache: &PipelineCache,
     skins_uniform: &SkinUniforms,
     weights_uniform: &MorphUniforms,
-    mesh_id_to_image: &HashMap<AssetId<Mesh>, MorphTargetImage>,
+    morph_target_images: &MorphTargetImages,
     morph_targets: &mut HashMap<AssetId<Mesh>, MeshBindGroupPair>,
 ) {
     let (skin, prev_skin) = (&skins_uniform.current_buffer, &skins_uniform.prev_buffer);
@@ -4054,12 +4056,11 @@ fn prepare_mesh_morph_target_bind_groups_for_phase_using_uniforms(
         .and_then(|descriptors_buffer| descriptors_buffer.buffer());
 
     for (id, gpu_mesh) in meshes.iter() {
-        if !gpu_mesh.has_morph_targets() {
-            continue;
-        }
-        let Some(morph_targets_image) = mesh_id_to_image.get(&id) else {
+        let Some(handle) = &gpu_mesh.morph_target_handle else {
             continue;
         };
+
+        let morph_targets_image = morph_target_images.get(handle);
         let targets = MorphTargetsResource::Texture(&morph_targets_image.texture_view);
         let bind_group_pair = if is_skinned(&gpu_mesh.layout) {
             MeshBindGroupPair {

crates/bevy_pbr/src/wireframe.rs

@@ -40,7 +40,8 @@ use bevy_render::{
     },
     prelude::*,
     render_asset::{
-        prepare_assets, PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssets,
+        prepare_assets, AlreadyTaken, PrepareAssetError, RenderAsset, RenderAssetPlugin,
+        RenderAssets,
     },
     render_phase::{
         AddRenderCommand, BinnedPhaseItem, BinnedRenderPhasePlugin, BinnedRenderPhaseType,
@@ -948,13 +949,21 @@ impl AsAssetId for Mesh3dWireframe {
 impl RenderAsset for RenderWireframeMaterial {
     type SourceAsset = WireframeMaterial;
     type Param = ();
+    type Extracted = Self::SourceAsset;
+
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
+        _previous_gpu_asset: Option<&Self>,
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        Some(Ok(source_asset.clone()))
+    }
 
     fn prepare_asset(
-        source_asset: Self::SourceAsset,
+        source_asset: Self::Extracted,
         _asset_id: AssetId<Self::SourceAsset>,
         _param: &mut SystemParamItem<Self::Param>,
         _previous_asset: Option<&Self>,
-    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
+    ) -> Result<Self, PrepareAssetError<Self::Extracted>> {
         Ok(RenderWireframeMaterial {
             color: source_asset.color.to_linear().to_f32_array(),
             line_width: source_asset.line_width,

crates/bevy_post_process/src/auto_exposure/compensation_curve.rs

@@ -1,9 +1,9 @@
-use bevy_asset::{prelude::*, RenderAssetUsages};
+use bevy_asset::prelude::*;
 use bevy_ecs::system::{lifetimeless::SRes, SystemParamItem};
 use bevy_math::{cubic_splines::CubicGenerator, FloatExt, Vec2};
 use bevy_reflect::prelude::*;
 use bevy_render::{
-    render_asset::{AssetExtractionError, RenderAsset},
+    render_asset::{AlreadyTaken, RenderAsset},
     render_resource::{
         Extent3d, ShaderType, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
         TextureView, UniformBuffer,
@@ -188,15 +188,13 @@ impl RenderAsset for GpuAutoExposureCompensationCurve {
     type SourceAsset = AutoExposureCompensationCurve;
     type Param = (SRes<RenderDevice>, SRes<RenderQueue>);
 
-    fn asset_usage(_: &Self::SourceAsset) -> RenderAssetUsages {
-        RenderAssetUsages::RENDER_WORLD
-    }
+    type Extracted = Self::SourceAsset;
 
-    fn take_gpu_data(
-        source: &mut Self::SourceAsset,
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
         _previous_gpu_asset: Option<&Self>,
-    ) -> Result<Self::SourceAsset, AssetExtractionError> {
-        Ok(source.clone())
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        Some(Ok(source_asset.clone()))
     }
 
     fn prepare_asset(

crates/bevy_render/Cargo.toml

@@ -128,6 +128,7 @@ indexmap = { version = "2" }
 bitflags = "2"
 itertools = "0.14"
 weak-table = "0.3"
+slotmap = "1.1.1"
 
 [dev-dependencies]
 proptest = "1"

crates/bevy_render/src/mesh/mod.rs

@@ -5,14 +5,14 @@ pub mod morph;
 #[cfg(feature = "morph")]
 use crate::GpuResourceAppExt;
 use crate::{
-    render_asset::{AssetExtractionError, PrepareAssetError, RenderAsset, RenderAssetPlugin},
+    render_asset::{AlreadyTaken, PrepareAssetError, RenderAsset, RenderAssetPlugin},
     renderer::{RenderDevice, RenderQueue},
     texture::GpuImage,
     RenderApp,
 };
 use allocator::MeshAllocatorPlugin;
 use bevy_app::{App, Plugin};
-use bevy_asset::{AssetId, RenderAssetUsages};
+use bevy_asset::AssetId;
 use bevy_camera::primitives::MeshAabb;
 use bevy_ecs::{
     prelude::*,
@@ -26,7 +26,7 @@ use glam::Vec3;
 use wgpu::IndexFormat;
 
 #[cfg(feature = "morph")]
-use crate::mesh::morph::RenderMorphTargetAllocator;
+use crate::mesh::morph::{MorphTargetImageHandle, RenderMorphTargetAllocator};
 
 /// Makes sure that [`Mesh`]es are extracted and prepared for the GPU.
 /// Does *not* add the [`Mesh`] as an asset. Use [`MeshPlugin`] for that.
@@ -57,7 +57,7 @@ impl Plugin for MeshRenderAssetPlugin {
 }
 
 /// The render world representation of a [`Mesh`].
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct RenderMesh {
     /// The number of vertices in the mesh.
     pub vertex_count: u32,
@@ -77,6 +77,9 @@ pub struct RenderMesh {
     /// Combined with [`RenderMesh::buffer_info`], this specifies the complete
     /// layout of the buffers associated with this mesh.
     pub layout: MeshVertexBufferLayoutRef,
+
+    #[cfg(feature = "morph")]
+    pub morph_target_handle: Option<MorphTargetImageHandle>,
 }
 
 impl RenderMesh {
@@ -135,21 +138,23 @@ impl RenderAsset for RenderMesh {
         SResMut<RenderMorphTargetAllocator>,
     );
 
-    #[inline]
-    fn asset_usage(mesh: &Self::SourceAsset) -> RenderAssetUsages {
-        mesh.asset_usage
-    }
+    type Extracted = Self::SourceAsset;
 
-    fn take_gpu_data(
-        source: &mut Self::SourceAsset,
+    fn extract(
+        source_asset: &mut Self::SourceAsset,
         _previous_gpu_asset: Option<&Self>,
-    ) -> Result<Self::SourceAsset, AssetExtractionError> {
-        source
-            .take_gpu_data()
-            .map_err(|_| AssetExtractionError::AlreadyExtracted)
+    ) -> Option<Result<Self::Extracted, AlreadyTaken>> {
+        source_asset
+            .asset_usage
+            .extract(
+                source_asset,
+                Self::SourceAsset::take_gpu_data,
+                |source_asset| source_asset.clone().take_gpu_data(), // This could be done more idiomatically.
+            )
+            .map(|result| result.map_err(|_| AlreadyTaken))
     }
 
-    fn byte_len(mesh: &Self::SourceAsset) -> Option<usize> {
+    fn byte_len(mesh: &Self::Extracted) -> Option<usize> {
         let mut vertex_size = 0;
         for attribute_data in mesh.attributes() {
             let vertex_format = attribute_data.0.format;
@@ -163,7 +168,7 @@ impl RenderAsset for RenderMesh {
 
     /// Converts the extracted mesh into a [`RenderMesh`].
     fn prepare_asset(
-        mesh: Self::SourceAsset,
+        mesh: Self::Extracted,
         _mesh_id: AssetId<Self::SourceAsset>,
         (
             _render_device,
@@ -172,7 +177,7 @@ impl RenderAsset for RenderMesh {
             _render_morph_targets_allocator,
         ): &mut SystemParamItem<Self::Param>,
         _: Option<&Self>,
-    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
+    ) -> Result<Self, PrepareAssetError<Self::Extracted>> {
         let (buffer_info, index_format) = match mesh.indices() {
             Some(indices) => (
                 RenderMeshBufferInfo::Indexed {
@@ -200,15 +205,14 @@ impl RenderAsset for RenderMesh {
 
         // Place the morph displacements in an image if necessary.
         #[cfg(feature = "morph")]
-        if let Some(morph_targets) = mesh.morph_targets() {
+        let morph_target_handle = mesh.morph_targets().and_then(|morph_targets| {
             _render_morph_targets_allocator.allocate(
                 _render_device,
                 _render_queue,
-                _mesh_id,
                 morph_targets,
                 mesh.count_vertices(),
-            );
-        }
+            )
+        });
 
         Ok(RenderMesh {
             vertex_count: mesh.count_vertices() as u32,
@@ -219,15 +223,20 @@ impl RenderAsset for RenderMesh {
             buffer_info,
             key_bits,
             layout: mesh_vertex_buffer_layout,
+            #[cfg(feature = "morph")]
+            morph_target_handle,
         })
     }
 
     fn unload_asset(
+        self,
         _mesh_id: AssetId<Self::SourceAsset>,
         (_, _, _, _render_morph_targets_allocator): &mut SystemParamItem<Self::Param>,
     ) {
         // Free the morph target images if necessary.
         #[cfg(feature = "morph")]
-        _render_morph_targets_allocator.free(_mesh_id);
+        if let Some(handle) = self.morph_target_handle {
+            _render_morph_targets_allocator.free(handle);
+        }
     }
 }