[Merged by Bors] - Directly extract joints into SkinnedMeshJoints

crates/bevy_pbr/src/render/mesh.rs

@@ -172,11 +172,6 @@ pub fn extract_meshes(
     commands.insert_or_spawn_batch(not_caster_commands);
 }
 
-#[derive(Resource, Debug, Default)]
-pub struct ExtractedJoints {
-    pub buffer: Vec<Mat4>,
-}
-
 #[derive(Component)]
 pub struct SkinnedMeshJoints {
     pub index: u32,
@@ -188,19 +183,22 @@ impl SkinnedMeshJoints {
         skin: &SkinnedMesh,
         inverse_bindposes: &Assets<SkinnedMeshInverseBindposes>,
         joints: &Query<&GlobalTransform>,
-        buffer: &mut Vec<Mat4>,
+        buffer: &mut BufferVec<Mat4>,
     ) -> Option<Self> {
         let inverse_bindposes = inverse_bindposes.get(&skin.inverse_bindposes)?;
-        let bindposes = inverse_bindposes.iter();
-        let skin_joints = skin.joints.iter();
         let start = buffer.len();
-        for (inverse_bindpose, joint) in bindposes.zip(skin_joints).take(MAX_JOINTS) {
-            if let Ok(joint) = joints.get(*joint) {
-                buffer.push(joint.affine() * *inverse_bindpose);
-            } else {
-                buffer.truncate(start);
-                return None;
-            }
+        let target = start + skin.joints.len().min(MAX_JOINTS);
+        buffer.extend(
+            joints
+                .iter_many(&skin.joints)
+                .zip(inverse_bindposes.iter())
+                .map(|(joint, bindpose)| joint.affine() * *bindpose),
+        );
+        // iter_many will skip any failed fetches. This will cause it to assign the wrong bones,
+        // so just bail by truncating to the start.
+        if buffer.len() != target {
+            buffer.truncate(start);
+            return None;
         }
 
         // Pad to 256 byte alignment
@@ -221,13 +219,13 @@ impl SkinnedMeshJoints {
 pub fn extract_skinned_meshes(
     mut commands: Commands,
     mut previous_len: Local<usize>,
-    mut previous_joint_len: Local<usize>,
+    mut uniform: ResMut<SkinnedMeshUniform>,
     query: Extract<Query<(Entity, &ComputedVisibility, &SkinnedMesh)>>,
     inverse_bindposes: Extract<Res<Assets<SkinnedMeshInverseBindposes>>>,
     joint_query: Extract<Query<&GlobalTransform>>,
 ) {
+    uniform.buffer.clear();
     let mut values = Vec::with_capacity(*previous_len);
-    let mut joints = Vec::with_capacity(*previous_joint_len);
     let mut last_start = 0;
 
     for (entity, computed_visibility, skin) in &query {
@@ -236,21 +234,19 @@ pub fn extract_skinned_meshes(
         }
         // PERF: This can be expensive, can we move this to prepare?
         if let Some(skinned_joints) =
-            SkinnedMeshJoints::build(skin, &inverse_bindposes, &joint_query, &mut joints)
+            SkinnedMeshJoints::build(skin, &inverse_bindposes, &joint_query, &mut uniform.buffer)
         {
             last_start = last_start.max(skinned_joints.index as usize);
             values.push((entity, skinned_joints.to_buffer_index()));
         }
     }
 
     // Pad out the buffer to ensure that there's enough space for bindings
-    while joints.len() - last_start < MAX_JOINTS {
-        joints.push(Mat4::ZERO);
+    while uniform.buffer.len() - last_start < MAX_JOINTS {
+        uniform.buffer.push(Mat4::ZERO);
     }
 
     *previous_len = values.len();
-    *previous_joint_len = joints.len();
-    commands.insert_resource(ExtractedJoints { buffer: joints });
     commands.insert_or_spawn_batch(values);
 }
 
@@ -779,20 +775,14 @@ impl Default for SkinnedMeshUniform {
 pub fn prepare_skinned_meshes(
     render_device: Res<RenderDevice>,
     render_queue: Res<RenderQueue>,
-    extracted_joints: Res<ExtractedJoints>,
     mut skinned_mesh_uniform: ResMut<SkinnedMeshUniform>,
 ) {
-    if extracted_joints.buffer.is_empty() {
+    if skinned_mesh_uniform.buffer.is_empty() {
         return;
     }
 
-    skinned_mesh_uniform.buffer.clear();
-    skinned_mesh_uniform
-        .buffer
-        .reserve(extracted_joints.buffer.len(), &render_device);
-    for joint in &extracted_joints.buffer {
-        skinned_mesh_uniform.buffer.push(*joint);
-    }
+    let len = skinned_mesh_uniform.buffer.len();
+    skinned_mesh_uniform.buffer.reserve(len, &render_device);
     skinned_mesh_uniform
         .buffer
         .write_buffer(&render_device, &render_queue);

crates/bevy_render/src/render_resource/buffer_vec.rs

@@ -131,7 +131,18 @@ impl<T: Pod> BufferVec<T> {
         }
     }
 
+    pub fn truncate(&mut self, len: usize) {
+        self.values.truncate(len);
+    }
+
     pub fn clear(&mut self) {
         self.values.clear();
     }
 }
+
+impl<T: Pod> Extend<T> for BufferVec<T> {
+    #[inline]
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
+        self.values.extend(iter);
+    }
+}