|
| 1 | +use bevy_ecs::{ |
| 2 | + change_detection::{DetectChanges, Ref}, |
| 3 | + component::{ComponentId, Tick}, |
| 4 | + entity::Entity, |
| 5 | + ptr::UnsafeCellDeref, |
| 6 | + storage::{SparseSets, Table, TableRow}, |
| 7 | +}; |
| 8 | + |
| 9 | +use crate::{zip_exact, TraitImplMeta, TraitImplRegistry, TraitQuery}; |
| 10 | + |
| 11 | +/// Read-access to all components implementing a trait for a given entity. |
| 12 | +/// |
| 13 | +/// This supports change detection and detection for added objects via |
| 14 | +/// |
| 15 | +/// - [`ReadTraits::iter_changed`] |
| 16 | +/// - [`ReadTraits::iter_added`] |
| 17 | +pub struct ReadTraits<'a, Trait: ?Sized + TraitQuery> { |
| 18 | + // Read-only access to the global trait registry. |
| 19 | + // Since no one outside of the crate can name the registry type, |
| 20 | + // we can be confident that no write accesses will conflict with this. |
| 21 | + pub(crate) registry: &'a TraitImplRegistry<Trait>, |
| 22 | + pub(crate) table: &'a Table, |
| 23 | + pub(crate) table_row: TableRow, |
| 24 | + /// This grants shared access to all sparse set components, |
| 25 | + /// but in practice we will only read the components specified in `self.registry`. |
| 26 | + /// The fetch impl registers read-access for all of these components, |
| 27 | + /// so there will be no runtime conflicts. |
| 28 | + pub(crate) sparse_sets: &'a SparseSets, |
| 29 | + pub(crate) last_run: Tick, |
| 30 | + pub(crate) this_run: Tick, |
| 31 | +} |
| 32 | + |
| 33 | +#[doc(hidden)] |
| 34 | +pub type CombinedReadTraitsIter<'a, Trait> = |
| 35 | + std::iter::Chain<ReadTableTraitsIter<'a, Trait>, ReadSparseTraitsIter<'a, Trait>>; |
| 36 | + |
| 37 | +#[doc(hidden)] |
| 38 | +pub struct ReadTableTraitsIter<'a, Trait: ?Sized> { |
| 39 | + // SAFETY: These two iterators must have equal length. |
| 40 | + pub(crate) components: std::slice::Iter<'a, ComponentId>, |
| 41 | + pub(crate) meta: std::slice::Iter<'a, TraitImplMeta<Trait>>, |
| 42 | + pub(crate) table_row: TableRow, |
| 43 | + // Grants shared access to the components corresponding to `components` in this table. |
| 44 | + // Not all components are guaranteed to exist in the table. |
| 45 | + pub(crate) table: &'a Table, |
| 46 | + pub(crate) last_run: Tick, |
| 47 | + pub(crate) this_run: Tick, |
| 48 | +} |
| 49 | + |
| 50 | +impl<'a, Trait: ?Sized + TraitQuery> Iterator for ReadTableTraitsIter<'a, Trait> { |
| 51 | + type Item = Ref<'a, Trait>; |
| 52 | + fn next(&mut self) -> Option<Self::Item> { |
| 53 | + // Iterate the remaining table components that are registered, |
| 54 | + // until we find one that exists in the table. |
| 55 | + let (column, meta) = unsafe { zip_exact(&mut self.components, &mut self.meta) } |
| 56 | + .find_map(|(&component, meta)| self.table.get_column(component).zip(Some(meta)))?; |
| 57 | + // SAFETY: We have shared access to the entire column. |
| 58 | + let ptr = unsafe { |
| 59 | + column |
| 60 | + .get_data_ptr() |
| 61 | + .byte_add(self.table_row.as_usize() * meta.size_bytes) |
| 62 | + }; |
| 63 | + let trait_object = unsafe { meta.dyn_ctor.cast(ptr) }; |
| 64 | + |
| 65 | + // SAFETY: we know that the `table_row` is a valid index. |
| 66 | + // Read access has been registered, so we can dereference it immutably. |
| 67 | + let added_tick = unsafe { column.get_added_tick_unchecked(self.table_row).deref() }; |
| 68 | + let changed_tick = unsafe { column.get_changed_tick_unchecked(self.table_row).deref() }; |
| 69 | + |
| 70 | + Some(Ref::new( |
| 71 | + trait_object, |
| 72 | + added_tick, |
| 73 | + changed_tick, |
| 74 | + self.last_run, |
| 75 | + self.this_run, |
| 76 | + )) |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +#[doc(hidden)] |
| 81 | +pub struct ReadSparseTraitsIter<'a, Trait: ?Sized> { |
| 82 | + // SAFETY: These two iterators must have equal length. |
| 83 | + pub(crate) components: std::slice::Iter<'a, ComponentId>, |
| 84 | + pub(crate) meta: std::slice::Iter<'a, TraitImplMeta<Trait>>, |
| 85 | + pub(crate) entity: Entity, |
| 86 | + // Grants shared access to the components corresponding to both `components` and `entity`. |
| 87 | + pub(crate) sparse_sets: &'a SparseSets, |
| 88 | + pub(crate) last_run: Tick, |
| 89 | + pub(crate) this_run: Tick, |
| 90 | +} |
| 91 | + |
| 92 | +impl<'a, Trait: ?Sized + TraitQuery> Iterator for ReadSparseTraitsIter<'a, Trait> { |
| 93 | + type Item = Ref<'a, Trait>; |
| 94 | + fn next(&mut self) -> Option<Self::Item> { |
| 95 | + // Iterate the remaining sparse set components that are registered, |
| 96 | + // until we find one that exists in the archetype. |
| 97 | + let ((ptr, ticks_ptr), meta) = unsafe { zip_exact(&mut self.components, &mut self.meta) } |
| 98 | + .find_map(|(&component, meta)| { |
| 99 | + self.sparse_sets |
| 100 | + .get(component) |
| 101 | + .and_then(|set| set.get_with_ticks(self.entity)) |
| 102 | + .zip(Some(meta)) |
| 103 | + })?; |
| 104 | + let trait_object = unsafe { meta.dyn_ctor.cast(ptr) }; |
| 105 | + let added_tick = unsafe { ticks_ptr.added.deref() }; |
| 106 | + let changed_tick = unsafe { ticks_ptr.changed.deref() }; |
| 107 | + Some(Ref::new( |
| 108 | + trait_object, |
| 109 | + added_tick, |
| 110 | + changed_tick, |
| 111 | + self.last_run, |
| 112 | + self.this_run, |
| 113 | + )) |
| 114 | + } |
| 115 | +} |
| 116 | + |
| 117 | +impl<'w, Trait: ?Sized + TraitQuery> IntoIterator for ReadTraits<'w, Trait> { |
| 118 | + type Item = Ref<'w, Trait>; |
| 119 | + type IntoIter = CombinedReadTraitsIter<'w, Trait>; |
| 120 | + #[inline] |
| 121 | + fn into_iter(self) -> Self::IntoIter { |
| 122 | + let table = ReadTableTraitsIter { |
| 123 | + components: self.registry.table_components.iter(), |
| 124 | + meta: self.registry.table_meta.iter(), |
| 125 | + table: self.table, |
| 126 | + table_row: self.table_row, |
| 127 | + last_run: self.last_run, |
| 128 | + this_run: self.this_run, |
| 129 | + }; |
| 130 | + let sparse = ReadSparseTraitsIter { |
| 131 | + components: self.registry.sparse_components.iter(), |
| 132 | + meta: self.registry.sparse_meta.iter(), |
| 133 | + entity: self.table.entities()[self.table_row.as_usize()], |
| 134 | + sparse_sets: self.sparse_sets, |
| 135 | + last_run: self.last_run, |
| 136 | + this_run: self.this_run, |
| 137 | + }; |
| 138 | + table.chain(sparse) |
| 139 | + } |
| 140 | +} |
| 141 | + |
| 142 | +impl<'w, Trait: ?Sized + TraitQuery> IntoIterator for &ReadTraits<'w, Trait> { |
| 143 | + type Item = Ref<'w, Trait>; |
| 144 | + type IntoIter = CombinedReadTraitsIter<'w, Trait>; |
| 145 | + #[inline] |
| 146 | + fn into_iter(self) -> Self::IntoIter { |
| 147 | + let table = ReadTableTraitsIter { |
| 148 | + components: self.registry.table_components.iter(), |
| 149 | + meta: self.registry.table_meta.iter(), |
| 150 | + table: self.table, |
| 151 | + table_row: self.table_row, |
| 152 | + last_run: self.last_run, |
| 153 | + this_run: self.this_run, |
| 154 | + }; |
| 155 | + let sparse = ReadSparseTraitsIter { |
| 156 | + components: self.registry.sparse_components.iter(), |
| 157 | + meta: self.registry.sparse_meta.iter(), |
| 158 | + entity: self.table.entities()[self.table_row.as_usize()], |
| 159 | + sparse_sets: self.sparse_sets, |
| 160 | + last_run: self.last_run, |
| 161 | + this_run: self.this_run, |
| 162 | + }; |
| 163 | + table.chain(sparse) |
| 164 | + } |
| 165 | +} |
| 166 | + |
| 167 | +impl<'w, Trait: ?Sized + TraitQuery> ReadTraits<'w, Trait> { |
| 168 | + /// Returns an iterator over the components implementing `Trait` for the current entity. |
| 169 | + pub fn iter(&self) -> CombinedReadTraitsIter<'w, Trait> { |
| 170 | + self.into_iter() |
| 171 | + } |
| 172 | + |
| 173 | + /// Returns an iterator over the components implementing `Trait` for the current entity |
| 174 | + /// that were added since the last time the system was run. |
| 175 | + pub fn iter_added(&self) -> impl Iterator<Item = Ref<'w, Trait>> { |
| 176 | + self.iter().filter(DetectChanges::is_added) |
| 177 | + } |
| 178 | + |
| 179 | + /// Returns an iterator over the components implementing `Trait` for the current entity |
| 180 | + /// whose values were changed since the last time the system was run. |
| 181 | + pub fn iter_changed(&self) -> impl Iterator<Item = Ref<'w, Trait>> { |
| 182 | + self.iter().filter(DetectChanges::is_changed) |
| 183 | + } |
| 184 | +} |
0 commit comments