Skip to content

Commit 859a3cb

Browse files
james7132viridiajanis-bhmkristoff3ralice-i-cecile
authored
Mitigate stack overflow on large Bundle inserts (#20772)
# Objective Fix #20571. ## Solution * Avoid passing the bundle by value further than one layer deep, and pass a `MovingPtr<'_, T>` of the Bundle instead. * Pass `MovingPtr<'_, Self>` to `DynamicBundle::get_components` and its recursive children. * Instead of returning an `BundleEffect`, directly apply the effect from a `MovingPtr<'_, MaybeUninit<Self>>`. * Remove the now unused `BundleEffect` trait. This should avoid most if not all extra stack copies of the bundle and its components. This won't 100% fix stack overflows via bundles, but it does mitigate them until much larger bundles are used. This started as a subset of the changes made in #20593. ## Testing Ran `cargo r --example feathers --features="experimental_bevy_feathers"` on Windows, no stack overflow. Co-Authored By: janis <janis@nirgendwo.xyz> --------- Co-authored-by: Talin <viridia@gmail.com> Co-authored-by: Janis <janis@nirgendwo.xyz> Co-authored-by: Janis <130913856+janis-bhm@users.noreply.github.com> Co-authored-by: Kristoffer Søholm <k.soeholm@gmail.com> Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com> Co-authored-by: Mike <mike.hsu@gmail.com> Co-authored-by: Chris Russell <8494645+chescock@users.noreply.github.com>
1 parent 8056c56 commit 859a3cb

File tree

16 files changed

+686
-312
lines changed

16 files changed

+686
-312
lines changed

‎crates/bevy_ecs/compile_fail/tests/ui/bundle_on_drop_impl.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
use bevy_ecs::prelude::*;
2+
3+
#[derive(Component, Debug)]
4+
pub struct A(usize);
5+
6+
// this should fail since destructuring T: Drop cannot be split.
7+
#[derive(Bundle, Debug)]
8+
//~^ E0509
9+
pub struct DropBundle {
10+
component_a: A,
11+
}
12+
13+
impl Drop for DropBundle {
14+
fn drop(&mut self) {
15+
// Just need the impl
16+
}
17+
}

‎crates/bevy_ecs/macros/src/lib.rs‎

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,20 +117,23 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
117117

118118
let mut active_field_types = Vec::new();
119119
let mut active_field_tokens = Vec::new();
120+
let mut active_field_alias: Vec<proc_macro2::TokenStream> = Vec::new();
120121
let mut inactive_field_tokens = Vec::new();
121122
for (((i, field_type), field_kind), field) in field_type
122123
.iter()
123124
.enumerate()
124125
.zip(field_kind.iter())
125126
.zip(field.iter())
126127
{
128+
let field_alias = format_ident!("field_{}", i).to_token_stream();
127129
let field_tokens = match field {
128130
Some(field) => field.to_token_stream(),
129131
None => Index::from(i).to_token_stream(),
130132
};
131133
match field_kind {
132134
BundleFieldKind::Component => {
133135
active_field_types.push(field_type);
136+
active_field_alias.push(field_alias);
134137
active_field_tokens.push(field_tokens);
135138
}
136139

@@ -165,16 +168,46 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
165168
};
166169

167170
let dynamic_bundle_impl = quote! {
168-
#[allow(deprecated)]
171+
#[doc(hidden)]
172+
#[expect(dead_code, reason = "This is a static assertion.")]
173+
impl #impl_generics #struct_name #ty_generics #where_clause {
174+
// Types that implement `Drop` cannot have their fields moved out. The implementation in
175+
// `get_componments` avoids this with pointers, so there needs to be a static assertion
176+
// that this is a sound thing to do. See https://doc.rust-lang.org/error_codes/E0509.html
177+
// for more information.
178+
fn check_no_bundle_drop(self) {
179+
// This has no effect, but we need to make sure the compiler doesn't optimize it out
180+
// black_box is used to do this
181+
#( core::hint::black_box(self.#active_field_tokens); )*
182+
}
183+
}
184+
169185
impl #impl_generics #ecs_path::bundle::DynamicBundle for #struct_name #ty_generics #where_clause {
170186
type Effect = ();
171187
#[allow(unused_variables)]
172188
#[inline]
173-
fn get_components(
174-
self,
189+
unsafe fn get_components(
190+
ptr: #ecs_path::ptr::MovingPtr<'_, Self>,
175191
func: &mut impl FnMut(#ecs_path::component::StorageType, #ecs_path::ptr::OwningPtr<'_>)
176192
) {
177-
#(<#active_field_types as #ecs_path::bundle::DynamicBundle>::get_components(self.#active_field_tokens, &mut *func);)*
193+
use #ecs_path::__macro_exports::DebugCheckedUnwrap;
194+
195+
#( let #active_field_alias = ptr.move_field(|ptr| &raw mut (*ptr).#active_field_tokens); )*
196+
core::mem::forget(ptr);
197+
#(
198+
<#active_field_types as #ecs_path::bundle::DynamicBundle>::get_components(
199+
#active_field_alias.try_into().debug_checked_unwrap(),
200+
func
201+
);
202+
)*
203+
}
204+
205+
#[allow(unused_variables)]
206+
#[inline]
207+
unsafe fn apply_effect(
208+
ptr: #ecs_path::ptr::MovingPtr<'_, core::mem::MaybeUninit<Self>>,
209+
func: &mut #ecs_path::world::EntityWorldMut<'_>,
210+
) {
178211
}
179212
}
180213
};

‎crates/bevy_ecs/src/bundle/impls.rs‎

Lines changed: 48 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
use core::any::TypeId;
22

3-
use bevy_ptr::OwningPtr;
4-
use variadics_please::all_tuples;
3+
use bevy_ptr::{MovingPtr, OwningPtr};
4+
use core::mem::MaybeUninit;
5+
use variadics_please::all_tuples_enumerated;
56

67
use crate::{
7-
bundle::{Bundle, BundleEffect, BundleFromComponents, DynamicBundle, NoBundleEffect},
8+
bundle::{Bundle, BundleFromComponents, DynamicBundle, NoBundleEffect},
89
component::{Component, ComponentId, Components, ComponentsRegistrator, StorageType},
10+
query::DebugCheckedUnwrap,
911
world::EntityWorldMut,
1012
};
1113

@@ -40,13 +42,19 @@ unsafe impl<C: Component> BundleFromComponents for C {
4042
impl<C: Component> DynamicBundle for C {
4143
type Effect = ();
4244
#[inline]
43-
fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) -> Self::Effect {
44-
OwningPtr::make(self, |ptr| func(C::STORAGE_TYPE, ptr));
45+
unsafe fn get_components(
46+
ptr: MovingPtr<'_, Self>,
47+
func: &mut impl FnMut(StorageType, OwningPtr<'_>),
48+
) -> Self::Effect {
49+
func(C::STORAGE_TYPE, OwningPtr::from(ptr));
4550
}
51+
52+
#[inline]
53+
unsafe fn apply_effect(_ptr: MovingPtr<'_, MaybeUninit<Self>>, _entity: &mut EntityWorldMut) {}
4654
}
4755

4856
macro_rules! tuple_impl {
49-
($(#[$meta:meta])* $($name: ident),*) => {
57+
($(#[$meta:meta])* $(($index:tt, $name: ident, $alias: ident)),*) => {
5058
#[expect(
5159
clippy::allow_attributes,
5260
reason = "This is a tuple-related macro; as such, the lints below may not always apply."
@@ -125,59 +133,52 @@ macro_rules! tuple_impl {
125133
reason = "Zero-length tuples will generate a function body equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case."
126134
)]
127135
#[inline(always)]
128-
fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) -> Self::Effect {
129-
#[allow(
130-
non_snake_case,
131-
reason = "The names of these variables are provided by the caller, not by us."
132-
)]
133-
let ($(mut $name,)*) = self;
134-
($(
135-
$name.get_components(&mut *func),
136-
)*)
136+
unsafe fn get_components(ptr: MovingPtr<'_, Self>, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) {
137+
// SAFETY:
138+
// - All of the `move_field` calls all fetch distinct and valid fields within `Self`.
139+
// - If a field is `NoBundleEffect`, it's `apply_effect` is a no-op
140+
// and cannot move any value out of an invalid instance after this call.
141+
// - If a field is `!NoBundleEffect`, it must be valid since a safe
142+
// implementation of `DynamicBundle` only moves the value out only
143+
// once between `get_components` and `apply_effect`.
144+
bevy_ptr::deconstruct_moving_ptr!(ptr => ($($index => $alias,)*));
145+
// SAFETY:
146+
// - If `ptr` is aligned, then field_ptr is aligned properly. Rust tuples cannot be `repr(packed)`.
147+
$( $name::get_components($alias.try_into().debug_checked_unwrap(), func); )*
137148
}
138-
}
139-
}
140-
}
141-
142-
all_tuples!(
143-
#[doc(fake_variadic)]
144-
tuple_impl,
145-
0,
146-
15,
147-
B
148-
);
149149

150-
macro_rules! after_effect_impl {
151-
($(#[$meta:meta])* $($after_effect: ident),*) => {
152-
#[expect(
153-
clippy::allow_attributes,
154-
reason = "This is a tuple-related macro; as such, the lints below may not always apply."
155-
)]
156-
$(#[$meta])*
157-
impl<$($after_effect: BundleEffect),*> BundleEffect for ($($after_effect,)*) {
158150
#[allow(
159151
clippy::unused_unit,
160-
reason = "Zero-length tuples will generate a function body equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case.")
161-
]
162-
fn apply(self, _entity: &mut EntityWorldMut) {
163-
#[allow(
164-
non_snake_case,
165-
reason = "The names of these variables are provided by the caller, not by us."
166-
)]
167-
let ($($after_effect,)*) = self;
168-
$($after_effect.apply(_entity);)*
152+
reason = "Zero-length tuples will generate a function body equivalent to `()`; however, this macro is meant for all applicable tuples, and as such it makes no sense to rewrite it just for that case."
153+
)]
154+
#[inline(always)]
155+
unsafe fn apply_effect(ptr: MovingPtr<'_, MaybeUninit<Self>>, entity: &mut EntityWorldMut) {
156+
// SAFETY:
157+
// - All of the `move_field` calls all fetch distinct and valid fields within `Self`.
158+
// - If a field is `NoBundleEffect`, it's `apply_effect` is a no-op
159+
// and cannot move any value out of an invalid instance.
160+
// - If a field is `!NoBundleEffect`, it must be valid since a safe
161+
// implementation of `DynamicBundle` only moves the value out only
162+
// once between `get_components` and `apply_effect`.
163+
bevy_ptr::deconstruct_moving_ptr!(ptr: MaybeUninit => (
164+
$($index => $alias,)*
165+
));
166+
// SAFETY:
167+
// - If `ptr` is aligned, then field_ptr is aligned properly. Rust tuples cannot be `repr(packed)`.
168+
$( $name::apply_effect($alias.try_into().debug_checked_unwrap(), entity); )*
169169
}
170170
}
171171

172172
$(#[$meta])*
173-
impl<$($after_effect: NoBundleEffect),*> NoBundleEffect for ($($after_effect,)*) { }
173+
impl<$($name: NoBundleEffect),*> NoBundleEffect for ($($name,)*) {}
174174
}
175175
}
176176

177-
all_tuples!(
177+
all_tuples_enumerated!(
178178
#[doc(fake_variadic)]
179-
after_effect_impl,
179+
tuple_impl,
180180
0,
181181
15,
182-
P
182+
B,
183+
field_
183184
);

‎crates/bevy_ecs/src/bundle/info.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use bevy_platform::{
33
collections::{HashMap, HashSet},
44
hash::FixedHasher,
55
};
6-
use bevy_ptr::OwningPtr;
6+
use bevy_ptr::{MovingPtr, OwningPtr};
77
use bevy_utils::TypeIdMap;
88
use core::{any::TypeId, ptr::NonNull};
99
use indexmap::{IndexMap, IndexSet};
@@ -228,8 +228,14 @@ impl BundleInfo {
228228
/// which removes the need to look up the [`ArchetypeAfterBundleInsert`](crate::archetype::ArchetypeAfterBundleInsert)
229229
/// in the archetype graph, which requires ownership of the entity's current archetype.
230230
///
231+
/// Regardless of how this is used, [`apply_effect`] must be called at most once on `bundle` after this function is
232+
/// called if `T::Effect: !NoBundleEffect` before returning to user-space safe code before returning to user-space safe code.
233+
/// This is currently only doable via use of [`MovingPtr::partial_move`].
234+
///
231235
/// `table` must be the "new" table for `entity`. `table_row` must have space allocated for the
232236
/// `entity`, `bundle` must match this [`BundleInfo`]'s type
237+
///
238+
/// [`apply_effect`]: crate::bundle::DynamicBundle::apply_effect
233239
#[inline]
234240
pub(super) unsafe fn write_components<'a, T: DynamicBundle, S: BundleComponentStatus>(
235241
&self,
@@ -240,14 +246,14 @@ impl BundleInfo {
240246
entity: Entity,
241247
table_row: TableRow,
242248
change_tick: Tick,
243-
bundle: T,
249+
bundle: MovingPtr<'_, T>,
244250
insert_mode: InsertMode,
245251
caller: MaybeLocation,
246-
) -> T::Effect {
252+
) {
247253
// NOTE: get_components calls this closure on each component in "bundle order".
248254
// bundle_info.component_ids are also in "bundle order"
249255
let mut bundle_component = 0;
250-
let after_effect = bundle.get_components(&mut |storage_type, component_ptr| {
256+
T::get_components(bundle, &mut |storage_type, component_ptr| {
251257
let component_id = *self
252258
.contributed_component_ids
253259
.get_unchecked(bundle_component);
@@ -303,8 +309,6 @@ impl BundleInfo {
303309
caller,
304310
);
305311
}
306-
307-
after_effect
308312
}
309313

310314
/// Internal method to initialize a required component from an [`OwningPtr`]. This should ultimately be called

‎crates/bevy_ecs/src/bundle/insert.rs‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use alloc::vec::Vec;
2-
use bevy_ptr::ConstNonNull;
2+
use bevy_ptr::{ConstNonNull, MovingPtr};
33
use core::ptr::NonNull;
44

55
use crate::{
@@ -138,18 +138,25 @@ impl<'w> BundleInserter<'w> {
138138
}
139139

140140
/// # Safety
141-
/// `entity` must currently exist in the source archetype for this inserter. `location`
142-
/// must be `entity`'s location in the archetype. `T` must match this [`BundleInfo`]'s type
141+
/// - `entity` must currently exist in the source archetype for this inserter.
142+
/// - `location` must be `entity`'s location in the archetype.
143+
/// - `T` must match this [`BundleInserter`] type used to create
144+
/// - If `T::Effect: !NoBundleEffect.`, then [`apply_effect`] must be called at most once on
145+
/// `bundle` after this function before returning to user-space safe code.
146+
/// - The value pointed to by `bundle` must not be accessed for anything other than [`apply_effect`]
147+
/// or dropped.
148+
///
149+
/// [`apply_effect`]: crate::bundle::DynamicBundle::apply_effect
143150
#[inline]
144151
pub(crate) unsafe fn insert<T: DynamicBundle>(
145152
&mut self,
146153
entity: Entity,
147154
location: EntityLocation,
148-
bundle: T,
155+
bundle: MovingPtr<'_, T>,
149156
insert_mode: InsertMode,
150157
caller: MaybeLocation,
151158
relationship_hook_mode: RelationshipHookMode,
152-
) -> (EntityLocation, T::Effect) {
159+
) -> EntityLocation {
153160
let bundle_info = self.bundle_info.as_ref();
154161
let archetype_after_insert = self.archetype_after_insert.as_ref();
155162
let archetype = self.archetype.as_ref();
@@ -188,15 +195,15 @@ impl<'w> BundleInserter<'w> {
188195
// so this reference can only be promoted from shared to &mut down here, after they have been ran
189196
let archetype = self.archetype.as_mut();
190197

191-
let (new_archetype, new_location, after_effect) = match &mut self.archetype_move_type {
198+
let (new_archetype, new_location) = match &mut self.archetype_move_type {
192199
ArchetypeMoveType::SameArchetype => {
193200
// SAFETY: Mutable references do not alias and will be dropped after this block
194201
let sparse_sets = {
195202
let world = self.world.world_mut();
196203
&mut world.storages.sparse_sets
197204
};
198205

199-
let after_effect = bundle_info.write_components(
206+
bundle_info.write_components(
200207
table,
201208
sparse_sets,
202209
archetype_after_insert,
@@ -209,7 +216,7 @@ impl<'w> BundleInserter<'w> {
209216
caller,
210217
);
211218

212-
(archetype, location, after_effect)
219+
(archetype, location)
213220
}
214221
ArchetypeMoveType::NewArchetypeSameTable { new_archetype } => {
215222
let new_archetype = new_archetype.as_mut();
@@ -237,7 +244,7 @@ impl<'w> BundleInserter<'w> {
237244
}
238245
let new_location = new_archetype.allocate(entity, result.table_row);
239246
entities.set(entity.index(), Some(new_location));
240-
let after_effect = bundle_info.write_components(
247+
bundle_info.write_components(
241248
table,
242249
sparse_sets,
243250
archetype_after_insert,
@@ -250,7 +257,7 @@ impl<'w> BundleInserter<'w> {
250257
caller,
251258
);
252259

253-
(new_archetype, new_location, after_effect)
260+
(new_archetype, new_location)
254261
}
255262
ArchetypeMoveType::NewArchetypeNewTable {
256263
new_archetype,
@@ -319,7 +326,7 @@ impl<'w> BundleInserter<'w> {
319326
}
320327
}
321328

322-
let after_effect = bundle_info.write_components(
329+
bundle_info.write_components(
323330
new_table,
324331
sparse_sets,
325332
archetype_after_insert,
@@ -332,7 +339,7 @@ impl<'w> BundleInserter<'w> {
332339
caller,
333340
);
334341

335-
(new_archetype, new_location, after_effect)
342+
(new_archetype, new_location)
336343
}
337344
};
338345

@@ -411,7 +418,7 @@ impl<'w> BundleInserter<'w> {
411418
}
412419
}
413420

414-
(new_location, after_effect)
421+
new_location
415422
}
416423

417424
#[inline]

0 commit comments

Comments
 (0)