Extract sprites into a Vec (#17619)

# Objective

Extract sprites into a `Vec` instead of a `HashMap`.

## Solution

Extract UI nodes into a `Vec` instead of an `EntityHashMap`.
Add an index into the `Vec` to `Transparent2d`.
Compare both the index and render entity in prepare so there aren't any
collisions.

## Showcase
yellow this PR, red main

```
cargo run --example many_sprites --release --features "trace_tracy"
```

`extract_sprites`
<img width="452" alt="extract_sprites"
src="https://github.com/user-attachments/assets/66c60406-7c2b-4367-907d-4a71d3630296"
/>

`queue_sprites`
<img width="463" alt="queue_sprites"
src="https://github.com/user-attachments/assets/54b903bd-4137-4772-9f87-e10e1e050d69"
/>

---------

Co-authored-by: Alice Cecile <[email protected]>

Author: ickshonpe

crates/bevy_core_pipeline/src/core_2d/mod.rs

@@ -349,6 +349,7 @@ pub struct Transparent2d {
     pub pipeline: CachedRenderPipelineId,
     pub draw_function: DrawFunctionId,
     pub batch_range: Range<u32>,
+    pub extracted_index: usize,
     pub extra_index: PhaseItemExtraIndex,
     /// Whether the mesh in question is indexed (uses an index buffer in
     /// addition to its vertex buffer).

crates/bevy_gizmos/src/pipeline_2d.rs

@@ -341,6 +341,7 @@ fn queue_line_gizmos_2d(
                     sort_key: FloatOrd(f32::INFINITY),
                     batch_range: 0..1,
                     extra_index: PhaseItemExtraIndex::None,
+                    extracted_index: usize::MAX,
                     indexed: false,
                 });
             }
@@ -362,6 +363,7 @@ fn queue_line_gizmos_2d(
                     sort_key: FloatOrd(f32::INFINITY),
                     batch_range: 0..1,
                     extra_index: PhaseItemExtraIndex::None,
+                    extracted_index: usize::MAX,
                     indexed: false,
                 });
             }
@@ -421,6 +423,7 @@ fn queue_line_joint_gizmos_2d(
                 sort_key: FloatOrd(f32::INFINITY),
                 batch_range: 0..1,
                 extra_index: PhaseItemExtraIndex::None,
+                extracted_index: usize::MAX,
                 indexed: false,
             });
         }

crates/bevy_sprite/src/mesh2d/material.rs

@@ -880,6 +880,7 @@ pub fn queue_material2d_meshes<M: Material2d>(
                         // Batching is done in batch_and_prepare_render_phase
                         batch_range: 0..1,
                         extra_index: PhaseItemExtraIndex::None,
+                        extracted_index: usize::MAX,
                         indexed: mesh.indexed(),
                     });
                 }

crates/bevy_sprite/src/render/mod.rs

@@ -19,7 +19,6 @@ use bevy_ecs::{
 use bevy_image::{BevyDefault, Image, ImageSampler, TextureAtlasLayout, TextureFormatPixelInfo};
 use bevy_math::{Affine3A, FloatOrd, Quat, Rect, Vec2, Vec4};
 use bevy_platform_support::collections::HashMap;
-use bevy_render::sync_world::MainEntity;
 use bevy_render::view::{RenderVisibleEntities, RetainedViewEntity};
 use bevy_render::{
     render_asset::RenderAssets,
@@ -32,7 +31,7 @@ use bevy_render::{
         *,
     },
     renderer::{RenderDevice, RenderQueue},
-    sync_world::{RenderEntity, TemporaryRenderEntity},
+    sync_world::RenderEntity,
     texture::{DefaultImageSampler, FallbackImage, GpuImage},
     view::{
         ExtractedView, Msaa, ViewTarget, ViewUniform, ViewUniformOffset, ViewUniforms,
@@ -339,13 +338,15 @@ pub struct ExtractedSprite {
     pub anchor: Vec2,
     /// For cases where additional [`ExtractedSprites`] are created during extraction, this stores the
     /// entity that caused that creation for use in determining visibility.
-    pub original_entity: Option<Entity>,
+    pub original_entity: Entity,
     pub scaling_mode: Option<ScalingMode>,
+    pub render_entity: Entity,
 }
 
 #[derive(Resource, Default)]
 pub struct ExtractedSprites {
-    pub sprites: HashMap<(Entity, MainEntity), ExtractedSprite>,
+    //pub sprites: HashMap<(Entity, MainEntity), ExtractedSprite>,
+    pub sprites: Vec<ExtractedSprite>,
 }
 
 #[derive(Resource, Default)]
@@ -388,19 +389,12 @@ pub fn extract_sprites(
         }
 
         if let Some(slices) = slices {
-            extracted_sprites.sprites.extend(
-                slices
-                    .extract_sprites(transform, original_entity, sprite)
-                    .map(|e| {
-                        (
-                            (
-                                commands.spawn(TemporaryRenderEntity).id(),
-                                original_entity.into(),
-                            ),
-                            e,
-                        )
-                    }),
-            );
+            extracted_sprites.sprites.extend(slices.extract_sprites(
+                &mut commands,
+                transform,
+                original_entity,
+                sprite,
+            ));
         } else {
             let atlas_rect = sprite
                 .texture_atlas
@@ -419,22 +413,20 @@ pub fn extract_sprites(
             };
 
             // PERF: we don't check in this function that the `Image` asset is ready, since it should be in most cases and hashing the handle is expensive
-            extracted_sprites.sprites.insert(
-                (entity, original_entity.into()),
-                ExtractedSprite {
-                    color: sprite.color.into(),
-                    transform: *transform,
-                    rect,
-                    // Pass the custom size
-                    custom_size: sprite.custom_size,
-                    flip_x: sprite.flip_x,
-                    flip_y: sprite.flip_y,
-                    image_handle_id: sprite.image.id(),
-                    anchor: sprite.anchor.as_vec(),
-                    original_entity: Some(original_entity),
-                    scaling_mode: sprite.image_mode.scale(),
-                },
-            );
+            extracted_sprites.sprites.push(ExtractedSprite {
+                render_entity: entity,
+                color: sprite.color.into(),
+                transform: *transform,
+                rect,
+                // Pass the custom size
+                custom_size: sprite.custom_size,
+                flip_x: sprite.flip_x,
+                flip_y: sprite.flip_y,
+                image_handle_id: sprite.image.id(),
+                anchor: sprite.anchor.as_vec(),
+                original_entity,
+                scaling_mode: sprite.image_mode.scale(),
+            });
         }
     }
 }
@@ -561,10 +553,10 @@ pub fn queue_sprites(
             .items
             .reserve(extracted_sprites.sprites.len());
 
-        for ((entity, main_entity), extracted_sprite) in extracted_sprites.sprites.iter() {
-            let index = extracted_sprite.original_entity.unwrap_or(*entity).index();
+        for (index, extracted_sprite) in extracted_sprites.sprites.iter().enumerate() {
+            let view_index = extracted_sprite.original_entity.index();
 
-            if !view_entities.contains(index as usize) {
+            if !view_entities.contains(view_index as usize) {
                 continue;
             }
 
@@ -575,11 +567,15 @@ pub fn queue_sprites(
             transparent_phase.add(Transparent2d {
                 draw_function: draw_sprite_function,
                 pipeline,
-                entity: (*entity, *main_entity),
+                entity: (
+                    extracted_sprite.render_entity,
+                    extracted_sprite.original_entity.into(),
+                ),
                 sort_key,
                 // `batch_range` is calculated in `prepare_sprite_image_bind_groups`
                 batch_range: 0..0,
                 extra_index: PhaseItemExtraIndex::None,
+                extracted_index: index,
                 indexed: true,
             });
         }
@@ -664,7 +660,12 @@ pub fn prepare_sprite_image_bind_groups(
         // Compatible items share the same entity.
         for item_index in 0..transparent_phase.items.len() {
             let item = &transparent_phase.items[item_index];
-            let Some(extracted_sprite) = extracted_sprites.sprites.get(&item.entity) else {
+
+            let Some(extracted_sprite) = extracted_sprites
+                .sprites
+                .get(item.extracted_index)
+                .filter(|extracted_sprite| extracted_sprite.render_entity == item.entity())
+            else {
                 // If there is a phase item that is not a sprite, then we must start a new
                 // batch to draw the other phase item(s) and to respect draw order. This can be
                 // done by invalidating the batch_image_handle

crates/bevy_sprite/src/texture_slice/computed_slices.rs

@@ -6,6 +6,7 @@ use bevy_ecs::prelude::*;
 use bevy_image::Image;
 use bevy_math::{Rect, Vec2};
 use bevy_platform_support::collections::HashSet;
+use bevy_render::sync_world::TemporaryRenderEntity;
 use bevy_transform::prelude::*;
 
 /// Component storing texture slices for tiled or sliced sprite entities
@@ -24,12 +25,13 @@ impl ComputedTextureSlices {
     /// * `sprite` - The sprite component
     /// * `handle` - The sprite texture handle
     #[must_use]
-    pub(crate) fn extract_sprites<'a>(
+    pub(crate) fn extract_sprites<'a, 'w, 's>(
         &'a self,
+        commands: &'a mut Commands<'w, 's>,
         transform: &'a GlobalTransform,
         original_entity: Entity,
         sprite: &'a Sprite,
-    ) -> impl ExactSizeIterator<Item = ExtractedSprite> + 'a {
+    ) -> impl ExactSizeIterator<Item = ExtractedSprite> + 'a + use<'a, 'w, 's> {
         let mut flip = Vec2::ONE;
         let [mut flip_x, mut flip_y] = [false; 2];
         if sprite.flip_x {
@@ -44,7 +46,8 @@ impl ComputedTextureSlices {
             let offset = (slice.offset * flip).extend(0.0);
             let transform = transform.mul_transform(Transform::from_translation(offset));
             ExtractedSprite {
-                original_entity: Some(original_entity),
+                render_entity: commands.spawn(TemporaryRenderEntity).id(),
+                original_entity,
                 color: sprite.color.into(),
                 transform,
                 rect: Some(slice.texture_rect),

crates/bevy_text/src/text2d.rs

@@ -204,24 +204,19 @@ pub fn extract_text2d_sprite(
             }
             let atlas = texture_atlases.get(&atlas_info.texture_atlas).unwrap();
 
-            extracted_sprites.sprites.insert(
-                (
-                    commands.spawn(TemporaryRenderEntity).id(),
-                    original_entity.into(),
-                ),
-                ExtractedSprite {
-                    transform: transform * GlobalTransform::from_translation(position.extend(0.)),
-                    color,
-                    rect: Some(atlas.textures[atlas_info.location.glyph_index].as_rect()),
-                    custom_size: None,
-                    image_handle_id: atlas_info.texture.id(),
-                    flip_x: false,
-                    flip_y: false,
-                    anchor: Anchor::Center.as_vec(),
-                    original_entity: Some(original_entity),
-                    scaling_mode: None,
-                },
-            );
+            extracted_sprites.sprites.push(ExtractedSprite {
+                render_entity: commands.spawn(TemporaryRenderEntity).id(),
+                transform: transform * GlobalTransform::from_translation(position.extend(0.)),
+                color,
+                rect: Some(atlas.textures[atlas_info.location.glyph_index].as_rect()),
+                custom_size: None,
+                image_handle_id: atlas_info.texture.id(),
+                flip_x: false,
+                flip_y: false,
+                anchor: Anchor::Center.as_vec(),
+                original_entity,
+                scaling_mode: None,
+            });
         }
     }
 }

examples/2d/mesh2d_manual.rs

@@ -412,6 +412,7 @@ pub fn queue_colored_mesh2d(
                     // This material is not batched
                     batch_range: 0..1,
                     extra_index: PhaseItemExtraIndex::None,
+                    extracted_index: usize::MAX,
                     indexed: mesh.indexed(),
                 });
             }