bevy_ecs/bundle/writer.rs
1use crate::{
2 component::{Component, ComponentId, Components, ComponentsRegistrator},
3 relationship::RelationshipHookMode,
4 world::EntityWorldMut,
5};
6use alloc::vec::Vec;
7use bevy_ptr::OwningPtr;
8use bumpalo::Bump;
9use core::{alloc::Layout, ptr::NonNull};
10
11/// Enables pushing components to internal scratch space (uses a bump allocator), which can then be
12/// written as a dynamic bundle. The contents are cleared after each write and the allocated scratch
13/// space is reused across writes.
14///
15/// Also see [`BundleWriter`].
16#[derive(Default)]
17pub struct BundleScratch {
18 // Correctness: this should never be made public or mismatched component ids could be inserted
19 component_ids: Vec<ComponentId>,
20 // Correctness: this should never be made public or arbitrary non-components could be inserted
21 component_ptrs: Vec<NonNull<u8>>,
22 // Safety: this cannot be exposed, otherwise `alloc.reset()` could be called in arbitrary places,
23 // which could invalidate the data stored in `component_ptrs`.
24 alloc: Bump,
25}
26
27impl BundleScratch {
28 /// Creates a new [`BundleWriter`] using this scratch space. For safety / correctness, this will
29 /// clear any existing components.
30 ///
31 /// Note that for performance reasons this will _not_ clear the internal allocator. To avoid leaking,
32 /// make sure every component pushed to the [`BundleWriter`] is followed by either a
33 /// [`BundleWriter::write`] or a [`BundleScratch::manual_drop`].
34 #[inline]
35 pub fn writer<'a>(&'a mut self) -> BundleWriter<'a> {
36 // This is necessary to ensure safety / correctness is maintained in the context of catch_unwind
37 // or a skipped `write`
38 self.component_ids.clear();
39 self.component_ptrs.clear();
40 BundleWriter(self)
41 }
42
43 /// Returns true if there are currently no components stored in the scratch space.
44 #[inline]
45 pub fn is_empty(&self) -> bool {
46 self.component_ids.is_empty()
47 }
48
49 /// This will drop all components currently stored in the scratch space. This is generally used to
50 /// ensure drops occur in error scenarios.
51 ///
52 /// # Safety
53 /// `components` must be from the same world as the components that were pushed to this writer.
54 pub unsafe fn manual_drop(&mut self, components: &Components) {
55 for (id, ptr) in self
56 .component_ids
57 .drain(..)
58 .zip(self.component_ptrs.drain(..))
59 {
60 if let Some(info) = components.get_info(id)
61 && let Some(drop) = info.drop()
62 {
63 // SAFETY: ptr is a valid component that matches the given component id
64 unsafe {
65 let ptr = OwningPtr::new(ptr);
66 (drop)(ptr);
67 }
68 }
69 }
70 self.alloc.reset();
71 }
72}
73
74/// Enables pushing components to the internal [`BundleScratch`], which can then be
75/// written as a dynamic bundle.
76///
77/// Components pushed to this writer should either be followed by a [`BundleWriter::write`] or a
78/// [`BundleScratch::manual_drop`] to avoid leaking.
79pub struct BundleWriter<'a>(&'a mut BundleScratch);
80
81// SAFETY: The `NonNull`s in component_ptrs are always a `Component`, which is Send
82unsafe impl Send for BundleScratch where Bump: Send {}
83
84impl<'a> BundleWriter<'a> {
85 /// Pushes the given component to the back of the current bundle scratch space. It will register
86 /// the component in `components` if it does not already exist.
87 ///
88 /// # Safety
89 ///
90 /// `components` must be from the same world that all previous [`Self::push_component`] or [`Self::push_component_by_id`] calls were called with,
91 /// and the _next_ [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] call.
92 pub unsafe fn push_component<C: Component>(
93 &mut self,
94 components: &mut ComponentsRegistrator,
95 component: C,
96 ) {
97 let id = components.register_component::<C>();
98 OwningPtr::make(component, |ptr| {
99 // SAFETY: ptr points to a C component value which matches the `id` looked up above.
100 // Layout matches C.
101 unsafe { self.push_component_by_id(id, ptr, Layout::new::<C>()) };
102 });
103 }
104
105 /// Pushes the given component ptr to the back of the current bundle scratch space.
106 ///
107 /// # Safety
108 ///
109 /// `components` must be from the same world that all previous [`Self::push_component`] or [`Self::push_component_by_id`] calls were called with,
110 /// and the _next_ [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] call. `component` must point to a [`Component`] value that matches `id`.
111 /// `layout` must correspond to the layout of the [`Component`] type.
112 pub unsafe fn push_component_by_id(
113 &mut self,
114 id: ComponentId,
115 component: OwningPtr<'_>,
116 layout: Layout,
117 ) {
118 let ptr = self.0.alloc.alloc_layout(layout);
119 // SAFETY:
120 // - `component` points to a valid value with this layout per precondition
121 // - `ptr` was just allocated (so cannot overlap) and has the same layout
122 unsafe { core::ptr::copy_nonoverlapping(component.as_ptr(), ptr.as_ptr(), layout.size()) };
123 self.0.component_ids.push(id);
124 self.0.component_ptrs.push(ptr);
125 }
126
127 /// Writes the current contents of the bundle to the given `entity` and clears the scratch space.
128 ///
129 /// Runs with [`RelationshipHookMode::Run`] by default.
130 /// Use [`write_with_relationship_hook_insert_mode`](Self::write_with_relationship_hook_insert_mode) if you need more flexibility.
131 ///
132 /// # Safety
133 ///
134 /// `entity` must be from the same world that all [`Self::push_component`] or [`Self::push_component_by_id`] calls since the last
135 /// [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] were called with.
136 #[track_caller]
137 pub unsafe fn write(self, entity: &mut EntityWorldMut) {
138 // SAFETY: Same preconditions
139 unsafe { self.write_with_relationship_hook_insert_mode(entity, RelationshipHookMode::Run) };
140 }
141
142 /// Writes the current contents of the bundle to the given `entity` and clears the scratch space.
143 ///
144 /// Also accepts [`RelationshipHookMode`] as an argument. Prefer [`write`](Self::write) to this if
145 /// [`RelationshipHookMode::Run`] by default is enough.
146 ///
147 /// # Safety
148 ///
149 /// `entity` must be from the same world that all [`Self::push_component`] or [`Self::push_component_by_id`] calls since the last
150 /// [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] were called with.
151 #[track_caller]
152 pub unsafe fn write_with_relationship_hook_insert_mode(
153 self,
154 entity: &mut EntityWorldMut,
155 relationship_hook_insert_mode: RelationshipHookMode,
156 ) {
157 // SAFETY:
158 // - All `component_ids` are from the same world as `entity`
159 // - All `component_data_ptrs` are valid types represented by `component_ids`
160 unsafe {
161 entity.insert_by_ids_internal(
162 &self.0.component_ids,
163 self.0
164 .component_ptrs
165 .drain(..)
166 .map(|ptr| OwningPtr::new(ptr)),
167 relationship_hook_insert_mode,
168 );
169 }
170 self.0.component_ids.clear();
171 self.0.alloc.reset();
172 }
173
174 /// Returns true if there are currently no components.
175 #[inline]
176 pub fn is_empty(&self) -> bool {
177 self.0.component_ids.is_empty()
178 }
179}
180
181#[cfg(test)]
182mod tests {
183 use crate::{bundle::BundleScratch, component::Component, name::Name, world::World};
184
185 #[test]
186 fn write_component() {
187 #[derive(Component)]
188 struct X;
189
190 let mut world = World::new();
191 let mut bundle_scratch = BundleScratch::default();
192 let mut bundle_writer = bundle_scratch.writer();
193 // SAFETY: the same world is used for every bundle_writer operation
194 unsafe {
195 let mut components = world.components_registrator();
196 bundle_writer.push_component(&mut components, X);
197 bundle_writer.push_component(&mut components, Name::new("Hi"));
198 let mut entity = world.spawn_empty();
199 bundle_writer.write(&mut entity);
200
201 assert_eq!(entity.get::<Name>().unwrap().as_str(), "Hi");
202 assert!(entity.contains::<X>());
203 }
204 }
205}