bevy_ecs/storage/table/

column.rs

use crate::{
    change_detection::{AtomicTick, CheckChangeTicks, ComponentTicks, MaybeLocation, Tick},
    component::ComponentInfo,
    storage::{blob_array::BlobArray, thin_array_ptr::ThinArrayPtr, TableRow},
};
use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
use core::{cell::UnsafeCell, mem::needs_drop, num::NonZeroUsize, panic::Location};

/// A type-erased contiguous container for data of a homogeneous type.
///
/// Conceptually, a `Column` is very similar to a type-erased `Box<[T]>`.
/// It also stores the change detection ticks for its components, kept in two separate
/// contiguous buffers internally. An element shares its data across these buffers by using the
/// same index (i.e. the entity at row 3 has it's data at index 3 and its change detection ticks at index 3).
///
/// Like many other low-level storage types, `Column` has a limited and highly unsafe
/// interface. It's highly advised to use higher level types and their safe abstractions
/// instead of working directly with `Column`.
///
/// For performance reasons, `Column` does not store its capacity and length.
/// This type is used by [`Table`] and [`ComponentSparseSet`], where the corresponding capacity
/// and length can be found.
///
/// [`Table`]: crate::storage::Table
/// [`ComponentSparseSet`]: crate::storage::ComponentSparseSet
#[derive(Debug)]
pub struct Column {
    data: BlobArray,
    added_ticks: ThinArrayPtr<UnsafeCell<Tick>>,
    changed_ticks: ThinArrayPtr<UnsafeCell<Tick>>,
    changed_by: MaybeLocation<ThinArrayPtr<UnsafeCell<&'static Location<'static>>>>,
    summary_tick: Option<AtomicTick>,
}

impl Column {
    /// Create a new [`Column`] with the given `capacity`.
    pub fn with_capacity(component_info: &ComponentInfo, capacity: usize) -> Self {
        Self {
            // SAFETY:
            // * The components stored in this columns will match the information in `component_info`
            // * `ComponentInfo` ensures that `layout().size()` is a multiple of `layout().align()`
            data: unsafe {
                BlobArray::with_capacity(component_info.layout(), component_info.drop(), capacity)
            },
            added_ticks: ThinArrayPtr::with_capacity(capacity),
            changed_ticks: ThinArrayPtr::with_capacity(capacity),
            changed_by: MaybeLocation::new_with(|| ThinArrayPtr::with_capacity(capacity)),
            summary_tick: if component_info.summary_tick() {
                // Set this to zero for now; when we initialize the column by
                // inserting a component it'll be updated with the correct
                // value.
                Some(AtomicTick::default())
            } else {
                None
            },
        }
    }

    /// Swap-remove and drop the removed element, but the component at `row` must not be the last element.
    ///
    /// # Safety
    /// - `row.as_usize()` < `len`
    /// - `last_element_index` = `len - 1`
    /// - `last_element_index` != `row.as_usize()`
    /// -   The caller should update the `len` to `len - 1`, or immediately initialize another element in the `last_element_index`
    pub(crate) unsafe fn swap_remove_and_drop_unchecked_nonoverlapping(
        &mut self,
        last_element_index: usize,
        row: TableRow,
    ) {
        self.data
            .swap_remove_and_drop_unchecked_nonoverlapping(row.index(), last_element_index);
        self.added_ticks
            .swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        self.changed_ticks
            .swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        self.changed_by.as_mut().map(|changed_by| {
            changed_by.swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        });
    }

    /// Swap-remove the provided row.
    ///
    /// If `DROP` is `true`, the removed element will be dropped as needed.
    ///
    /// If `DROP` is `false`, the removed element will be forgotten.
    ///
    /// # Safety
    /// - `last_element_index` must be the index of the last element.
    /// - `row.index()` <= `last_element_index`
    /// - The caller should update their saved length to reflect the change (decrement it by 1).
    pub(crate) unsafe fn swap_remove_unchecked<const DROP: bool>(
        &mut self,
        last_element_index: usize,
        row: TableRow,
    ) {
        if DROP {
            self.data
                .swap_remove_and_drop_unchecked(row.index(), last_element_index);
        } else {
            _ = self
                .data
                .swap_remove_unchecked(row.index(), last_element_index);
        }
        self.added_ticks
            .swap_remove_unchecked(row.index(), last_element_index);
        self.changed_ticks
            .swap_remove_unchecked(row.index(), last_element_index);
        self.changed_by.as_mut().map(|changed_by| {
            changed_by.swap_remove_unchecked(row.index(), last_element_index);
        });
    }

    /// Swap-remove and forgets the removed element, but the component at `row` must not be the last element.
    ///
    /// # Safety
    /// - `row.as_usize()` < `len`
    /// - `last_element_index` = `len - 1`
    /// - `last_element_index` != `row.as_usize()`
    /// -   The caller should update the `len` to `len - 1`, or immediately initialize another element in the `last_element_index`
    pub(crate) unsafe fn swap_remove_and_forget_unchecked_nonoverlapping(
        &mut self,
        last_element_index: usize,
        row: TableRow,
    ) -> OwningPtr<'_> {
        let data = self
            .data
            .swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        self.added_ticks
            .swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        self.changed_ticks
            .swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        self.changed_by.as_mut().map(|changed_by| {
            changed_by.swap_remove_unchecked_nonoverlapping(row.index(), last_element_index);
        });
        data
    }

    /// Call [`realloc`](std::alloc::realloc) to expand / shrink the memory allocation for this [`Column`]
    ///
    /// # Panics
    /// - Panics if the any of the new capacity overflows `isize::MAX` bytes.
    /// - Panics if the any of the reallocations causes an out-of-memory error.
    ///
    /// # Safety
    /// - `current_capacity` must be the current capacity of this column (the capacity of `self.data`, `self.added_ticks`, `self.changed_tick`)
    /// -   The caller should make sure their saved `capacity` value is updated to `new_capacity` after this operation.
    pub(crate) unsafe fn realloc(
        &mut self,
        current_capacity: NonZeroUsize,
        new_capacity: NonZeroUsize,
    ) {
        self.data.realloc(current_capacity, new_capacity);
        self.added_ticks.realloc(current_capacity, new_capacity);
        self.changed_ticks.realloc(current_capacity, new_capacity);
        self.changed_by
            .as_mut()
            .map(|changed_by| changed_by.realloc(current_capacity, new_capacity));
    }

    /// Call [`alloc`](std::alloc::alloc) to allocate memory for this [`Column`]
    /// The caller should make sure their saved `capacity` value is updated to `new_capacity` after this operation.
    ///
    /// # Panics
    /// - Panics if the any of the new capacity overflows `isize::MAX` bytes.
    /// - Panics if the any of the allocations causes an out-of-memory error.
    pub(crate) fn alloc(&mut self, new_capacity: NonZeroUsize) {
        self.data.alloc(new_capacity);
        self.added_ticks.alloc(new_capacity);
        self.changed_ticks.alloc(new_capacity);
        self.changed_by
            .as_mut()
            .map(|changed_by| changed_by.alloc(new_capacity));
    }

    /// Writes component data to the column at the given row.
    /// Assumes the slot is uninitialized, drop is not called.
    /// To overwrite existing initialized value, use [`Self::replace`] instead.
    ///
    /// # Safety
    /// - `row.as_usize()` must be in bounds.
    /// - `data` holds a component that matches the `component_id`
    #[inline]
    pub(crate) unsafe fn initialize(
        &mut self,
        row: TableRow,
        data: OwningPtr<'_>,
        tick: Tick,
        caller: MaybeLocation,
    ) {
        self.data.initialize_unchecked(row.index(), data);
        self.added_ticks
            .initialize_unchecked(row.index(), UnsafeCell::new(tick));
        self.changed_ticks
            .initialize_unchecked(row.index(), UnsafeCell::new(tick));
        self.changed_by
            .as_mut()
            .zip(caller)
            .map(|(changed_by, caller)| {
                changed_by.initialize_unchecked(row.index(), UnsafeCell::new(caller));
            });
        if let Some(summary_tick) = &self.summary_tick {
            Self::update_summary_tick(summary_tick, tick);
        }
    }

    /// Overwrites component data to the column at given row. The previous value is dropped.
    ///
    /// # Safety
    /// - There must be a valid initialized value stored at `row`.
    /// - `row.as_usize()` must be in bounds.
    /// - `data` holds a component that matches the `component_id`
    #[inline]
    pub(crate) unsafe fn replace(
        &mut self,
        row: TableRow,
        data: OwningPtr<'_>,
        change_tick: Tick,
        caller: MaybeLocation,
    ) {
        self.data.replace_unchecked(row.index(), data);
        *self.changed_ticks.get_unchecked_mut(row.index()).get_mut() = change_tick;
        self.changed_by
            .as_mut()
            .map(|changed_by| changed_by.get_unchecked_mut(row.index()).get_mut())
            .assign(caller);
        if let Some(summary_tick) = &self.summary_tick {
            Self::update_summary_tick(summary_tick, change_tick);
        }
    }

    /// Removes the element from `src` at `src_row` and inserts it
    /// into this column to initialize the values at `dst_row`.
    /// Does not do any bounds checking.
    ///
    /// `change_tick` must be the change tick of the current system.
    ///
    /// # Safety
    ///  - `src` must have the same data layout as `self`
    ///  - `src_row` must be in bounds for `src`
    ///  - `dst_row` must be in bounds for `self`
    ///  - `src[src_row]` must be initialized to a valid value.
    ///  - `self[dst_row]` must not be initialized yet.
    ///  - `src_last_element_index` must be the last element in `src`
    #[inline]
    pub(crate) unsafe fn initialize_from_unchecked(
        &mut self,
        src: &mut Column,
        src_last_element_index: usize,
        src_row: TableRow,
        dst_row: TableRow,
        this_run: Tick,
    ) {
        debug_assert!(self.data.layout() == src.data.layout());

        // Making this a cold path avoids most of the performance cost.
        #[cold]
        fn update_summary_tick_from_row(
            changed_ticks: &ThinArrayPtr<UnsafeCell<Tick>>,
            summary_tick: &AtomicTick,
            dst_row: TableRow,
            this_run: Tick,
        ) {
            // SAFETY:
            // - Changed tick just got initialized at dst_row
            // - There are no mutable references to the changed tick
            let row_change_tick = unsafe { changed_ticks.get_unchecked(dst_row.index()).read() };
            if row_change_tick.is_newer_than(summary_tick.get(), this_run) {
                summary_tick.set(row_change_tick);
            }
        }

        // SAFETY:
        // In bounds, same layout & correct last element index as per preconditions
        unsafe {
            self.data.initialize_from_swap_remove_unchecked(
                &mut src.data,
                src_last_element_index,
                src_row.index(),
                dst_row.index(),
            );
            self.added_ticks.initialize_from_swap_remove_unchecked(
                &mut src.added_ticks,
                src_last_element_index,
                src_row.index(),
                dst_row.index(),
            );
            self.changed_ticks.initialize_from_swap_remove_unchecked(
                &mut src.changed_ticks,
                src_last_element_index,
                src_row.index(),
                dst_row.index(),
            );
            self.changed_by.as_mut().zip(src.changed_by.as_mut()).map(
                |(self_changed_by, src_changed_by)| {
                    self_changed_by.initialize_from_swap_remove_unchecked(
                        src_changed_by,
                        src_last_element_index,
                        src_row.index(),
                        dst_row.index(),
                    );
                },
            );
        }

        if let Some(summary_tick) = &self.summary_tick {
            update_summary_tick_from_row(&self.changed_ticks, summary_tick, dst_row, this_run);
        }
    }

    /// Calls [`Tick::check_tick`] on all of the ticks stored in this column, as
    /// well as the summary tick, if one is present.
    ///
    /// # Safety
    /// `len` is the actual length of this column
    #[inline]
    pub(crate) unsafe fn check_change_ticks(&mut self, len: usize, check: CheckChangeTicks) {
        for i in 0..len {
            // SAFETY:
            // - `i` < `len`
            // we have a mutable reference to `self`
            unsafe { self.added_ticks.get_unchecked_mut(i) }
                .get_mut()
                .check_tick(check);
            // SAFETY:
            // - `i` < `len`
            // we have a mutable reference to `self`
            unsafe { self.changed_ticks.get_unchecked_mut(i) }
                .get_mut()
                .check_tick(check);
        }

        // Update the summary tick, if one is present.
        if let Some(summary_tick) = &self.summary_tick {
            let mut summary_tick_value = summary_tick.get();
            if summary_tick_value.check_tick(check) {
                summary_tick.set(summary_tick_value);
            }
        }
    }

    /// Clear all the components from this column.
    ///
    /// # Safety
    /// - `len` must match the actual length of the column
    /// -   The caller must not use the elements this column's data until [`initializing`](Self::initialize) it again (set `len` to 0).
    pub(crate) unsafe fn clear(&mut self, len: usize) {
        self.added_ticks.clear_elements(len);
        self.changed_ticks.clear_elements(len);
        self.data.clear(len);
        self.changed_by
            .as_mut()
            .map(|changed_by| changed_by.clear_elements(len));
    }

    /// Because this method needs parameters, it can't be the implementation of the `Drop` trait.
    /// The owner of this [`Column`] must call this method with the correct information.
    ///
    /// # Safety
    /// - `len` is indeed the length of the column
    /// - `cap` is indeed the capacity of the column
    /// - the data stored in `self` will never be used again
    pub(crate) unsafe fn drop(&mut self, cap: usize, len: usize) {
        self.added_ticks.drop(cap, len);
        self.changed_ticks.drop(cap, len);
        self.data.drop(cap, len);
        self.changed_by
            .as_mut()
            .map(|changed_by| changed_by.drop(cap, len));
    }

    /// Drops the last component in this column.
    ///
    /// # Safety
    /// - `last_element_index` is indeed the index of the last element
    /// - the data stored in `last_element_index` will never be used unless properly initialized again.
    pub(crate) unsafe fn drop_last_component(&mut self, last_element_index: usize) {
        const {
            assert!(!needs_drop::<UnsafeCell<Tick>>());
            assert!(!needs_drop::<UnsafeCell<&'static Location<'static>>>());
        }
        self.data.drop_last_element(last_element_index);
    }

    /// Get a slice to the data stored in this [`Column`].
    ///
    /// # Safety
    /// - `T` must match the type of data that's stored in this [`Column`]
    /// - `len` must match the actual length of this column (number of elements stored)
    pub unsafe fn get_data_slice<T>(&self, len: usize) -> &[UnsafeCell<T>] {
        // SAFETY: Upheld by caller
        unsafe { self.data.get_sub_slice(len) }
    }

    /// Get a slice to the added [`ticks`](Tick) in this [`Column`].
    ///
    /// # Safety
    /// - `len` must match the actual length of this column (number of elements stored)
    pub unsafe fn get_added_ticks_slice(&self, len: usize) -> &[UnsafeCell<Tick>] {
        // SAFETY: Upheld by caller
        unsafe { self.added_ticks.as_slice(len) }
    }

    /// Get a slice to the changed [`ticks`](Tick) in this [`Column`].
    ///
    /// # Safety
    /// - `len` must match the actual length of this column (number of elements stored)
    pub unsafe fn get_changed_ticks_slice(&self, len: usize) -> &[UnsafeCell<Tick>] {
        // SAFETY: Upheld by caller
        unsafe { self.changed_ticks.as_slice(len) }
    }

    /// Get a slice to the calling locations that last changed each value in this [`Column`]
    ///
    /// # Safety
    /// - `len` must match the actual length of this column (number of elements stored)
    pub unsafe fn get_changed_by_slice(
        &self,
        len: usize,
    ) -> MaybeLocation<&[UnsafeCell<&'static Location<'static>>]> {
        self.changed_by
            .as_ref()
            .map(|changed_by| changed_by.as_slice(len))
    }

    /// Fetches a read-only reference to the data at `row`. This does not
    /// do any bounds checking.
    ///
    /// # Safety
    /// - `row` must be within the range `[0, self.len())`.
    /// - no other mutable reference to the data of the same row can exist at the same time
    #[inline]
    pub unsafe fn get_data_unchecked(&self, row: TableRow) -> Ptr<'_> {
        self.data.get_unchecked(row.index())
    }

    /// Fetches the calling location that last changed the value at `row`.
    ///
    /// This function does not do any bounds checking.
    ///
    /// # Safety
    /// `row` must be within the range `[0, self.len())`.
    #[inline]
    pub unsafe fn get_changed_by_unchecked(
        &self,
        row: TableRow,
    ) -> MaybeLocation<&UnsafeCell<&'static Location<'static>>> {
        self.changed_by
            .as_ref()
            .map(|changed_by| changed_by.get_unchecked(row.index()))
    }

    /// Fetches the "added" change detection tick for the value at `row`.
    /// This function does not do any bounds checking.
    ///
    /// # Safety
    /// `row` must be within the range `[0, self.len())`.
    #[inline]
    pub unsafe fn get_added_tick_unchecked(&self, row: TableRow) -> &UnsafeCell<Tick> {
        // SAFETY: Upheld by caller
        unsafe { self.added_ticks.get_unchecked(row.index()) }
    }

    /// Fetches the "changed" change detection tick for the value at `row`
    /// This function does not do any bounds checking.
    ///
    /// # Safety
    /// `row` must be within the range `[0, self.len())`.
    #[inline]
    pub unsafe fn get_changed_tick_unchecked(&self, row: TableRow) -> &UnsafeCell<Tick> {
        // SAFETY: Upheld by caller
        unsafe { self.changed_ticks.get_unchecked(row.index()) }
    }

    /// Fetches the change detection ticks for the value at `row`.
    /// This function does not do any bounds checking.
    ///
    /// # Safety
    /// `row` must be within the range `[0, self.len())`.
    #[inline]
    pub unsafe fn get_ticks_unchecked(&self, row: TableRow) -> ComponentTicks {
        ComponentTicks {
            added: self.added_ticks.get_unchecked(row.index()).read(),
            changed: self.changed_ticks.get_unchecked(row.index()).read(),
        }
    }

    /// Returns the drop function for elements of the column,
    /// or `None` if they don't need to be dropped.
    #[inline]
    pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {
        self.data.get_drop()
    }

    /// Returns a reference to the summary tick for this column, if the
    /// component that this column is associated with has a summary tick.
    ///
    /// The summary tick stores the most recent changed timestamp that was
    /// written to any component instance of this column. "Most recent" here
    /// refers to the wall clock.
    ///
    /// Be careful when using this value, as it's easy to misuse. Because
    /// multiple systems with different ticks can be concurrently writing to a
    /// single column, and because "most recent" is in reference to wall clock
    /// time, *there may be ticks in the column that are logically later than
    /// the summary tick.* It's therefore incorrect to assume that the summary
    /// tick represents the latest tick stored in the column.
    ///
    /// Importantly, however, this situation can only occur when multiple
    /// systems are *actually* concurrently writing to a column. This can only
    /// happen when both systems are writing to a column with sparse queries.
    /// Systems that iterate over components with dense iteration have immutable
    /// or exclusive access to the columns corresponding to those components in
    /// the tables that they iterate over (and this fact is what makes
    /// contiguous iteration safe to begin with). So, *for a dense query*, we
    /// can guarantee that if the `last_run` tick for that query is *t*, then if
    /// the summary tick has a value earlier than *t*, then there have been no
    /// changes to the column values. Again, this is *only* true for dense
    /// iteration.
    #[inline]
    pub fn get_summary_tick(&self) -> Option<&AtomicTick> {
        self.summary_tick.as_ref()
    }

    /// Separate method for the sake of using `#[cold]`,
    /// since most components won't have summary ticks.
    #[cold]
    fn update_summary_tick(summary_tick: &AtomicTick, tick: Tick) {
        summary_tick.set(tick);
    }
}