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 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 core::ptr::copy(component.as_ptr(), ptr.as_ptr(), layout.size());
120 self.0.component_ids.push(id);
121 self.0.component_ptrs.push(ptr);
122 }
123
124 /// Writes the current contents of the bundle to the given `entity` and clears the scratch space.
125 ///
126 /// Runs with [`RelationshipHookMode::Run`] by default.
127 /// Use [`write_with_relationship_hook_insert_mode`](Self::write_with_relationship_hook_insert_mode) if you need more flexibility.
128 ///
129 /// # Safety
130 ///
131 /// `entity` must be from the same world that all [`Self::push_component`] or [`Self::push_component_by_id`] calls since the last
132 /// [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] were called with.
133 #[track_caller]
134 pub unsafe fn write(self, entity: &mut EntityWorldMut) {
135 self.write_with_relationship_hook_insert_mode(entity, RelationshipHookMode::Run);
136 }
137
138 /// Writes the current contents of the bundle to the given `entity` and clears the scratch space.
139 ///
140 /// Also accepts [`RelationshipHookMode`] as an argument. Prefer [`write`](Self::write) to this if
141 /// [`RelationshipHookMode::Run`] by default is enough.
142 ///
143 /// # Safety
144 ///
145 /// `entity` must be from the same world that all [`Self::push_component`] or [`Self::push_component_by_id`] calls since the last
146 /// [`Self::write`] or [`Self::write_with_relationship_hook_insert_mode`] were called with.
147 #[track_caller]
148 pub unsafe fn write_with_relationship_hook_insert_mode(
149 self,
150 entity: &mut EntityWorldMut,
151 relationship_hook_insert_mode: RelationshipHookMode,
152 ) {
153 // SAFETY:
154 // - All `component_ids` are from the same world as `entity`
155 // - All `component_data_ptrs` are valid types represented by `component_ids`
156 unsafe {
157 entity.insert_by_ids_internal(
158 &self.0.component_ids,
159 self.0
160 .component_ptrs
161 .drain(..)
162 .map(|ptr| OwningPtr::new(ptr)),
163 relationship_hook_insert_mode,
164 );
165 }
166 self.0.component_ids.clear();
167 self.0.alloc.reset();
168 }
169
170 /// Returns true if there are currently no components.
171 #[inline]
172 pub fn is_empty(&self) -> bool {
173 self.0.component_ids.is_empty()
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use crate::{bundle::BundleScratch, component::Component, name::Name, world::World};
180
181 #[test]
182 fn write_component() {
183 #[derive(Component)]
184 struct X;
185
186 let mut world = World::new();
187 let mut bundle_scratch = BundleScratch::default();
188 let mut bundle_writer = bundle_scratch.writer();
189 // SAFETY: the same world is used for every bundle_writer operation
190 unsafe {
191 let mut components = world.components_registrator();
192 bundle_writer.push_component(&mut components, X);
193 bundle_writer.push_component(&mut components, Name::new("Hi"));
194 let mut entity = world.spawn_empty();
195 bundle_writer.write(&mut entity);
196
197 assert_eq!(entity.get::<Name>().unwrap().as_str(), "Hi");
198 assert!(entity.contains::<X>());
199 }
200 }
201}