bevy_ecs 0.20.0

use alloc::vec::Vec;
use bevy_ptr::{ConstNonNull, MovingPtr};
use core::ptr::NonNull;

use crate::{
    archetype::{
        Archetype, ArchetypeAfterBundleInsert, ArchetypeCreated, ArchetypeId, Archetypes,
        ComponentStatus, ARCHETYPE_CREATED,
    },
    bundle::{ArchetypeMoveType, Bundle, BundleId, BundleInfo, DynamicBundle, InsertMode},
    change_detection::{MaybeLocation, Tick},
    component::{Components, StorageType},
    entity::{Entities, Entity, EntityLocation},
    event::{EntityComponentsTrigger, GlobalTrigger},
    lifecycle::{AddEvent, DiscardEvent, InsertEvent, ADD, DISCARD, INSERT},
    observer::Observers,
    query::DebugCheckedUnwrap as _,
    relationship::RelationshipHookMode,
    storage::{SparseSets, Storages, Table},
    world::{unsafe_world_cell::UnsafeWorldCell, World},
};

// SAFETY: We have exclusive world access so our pointers can't be invalidated externally
pub(crate) struct BundleInserter<'w> {
    /// Has complete read+write permissions
    world: UnsafeWorldCell<'w>,
    bundle_info: ConstNonNull<BundleInfo>,
    archetype_after_insert: ConstNonNull<ArchetypeAfterBundleInsert>,
    archetype: NonNull<Archetype>,
    archetype_move_type: ArchetypeMoveType,
    change_tick: Tick,
}

impl<'w> BundleInserter<'w> {
    /// # Safety
    /// - `archetype_id` must correspond to a valid archetype in `world`.
    #[inline]
    pub(crate) unsafe fn new<T: Bundle>(
        world: &'w mut World,
        archetype_id: ArchetypeId,
        change_tick: Tick,
    ) -> Self {
        let bundle_id = world.register_bundle_info::<T>();

        // SAFETY: We just ensured this bundle exists
        unsafe { Self::new_with_id(world, archetype_id, bundle_id, change_tick) }
    }

    /// Creates a new [`BundleInserter`].
    ///
    /// # Safety
    /// - `bundle_id` must correspond to an existing bundle in `world`.
    /// - `archetype_id` must correspond to a valid archetype in `world`.
    #[inline]
    pub(crate) unsafe fn new_with_id(
        world: &'w mut World,
        archetype_id: ArchetypeId,
        bundle_id: BundleId,
        change_tick: Tick,
    ) -> Self {
        // SAFETY: bundle exists per precondition
        let bundle_info = unsafe { world.bundles.get_unchecked(bundle_id) };
        // SAFETY: same world
        let (new_archetype_id, is_new_created) = unsafe {
            bundle_info.insert_bundle_into_archetype(
                &mut world.archetypes,
                &mut world.storages,
                &world.components,
                &world.observers,
                archetype_id,
            )
        };

        // SAFETY:
        // - The caller ensures `archetype_id` is valid.
        // - `new_archetype_id` was just created or fetched from the archetype graph.
        let (archetype, new_archetype) = unsafe {
            world
                .archetypes
                .get_maybe_disjoint_mut(archetype_id, new_archetype_id)
        };

        // SAFETY: The edge is assured to be initialized when we called insert_bundle_into_archetype
        let archetype_after_insert = unsafe {
            archetype
                .edges()
                .get_archetype_after_bundle_insert_internal(bundle_id)
                .debug_checked_unwrap()
                .into()
        };

        let archetype_move_type = if let Some(new_archetype) = new_archetype {
            if archetype.table_id() == new_archetype.table_id() {
                ArchetypeMoveType::NewArchetypeSameTable {
                    new_archetype: new_archetype.into(),
                }
            } else {
                ArchetypeMoveType::NewArchetypeNewTable {
                    new_archetype: new_archetype.into(),
                }
            }
        } else {
            ArchetypeMoveType::SameArchetype
        };

        let inserter = Self {
            archetype: archetype.into(),
            archetype_after_insert,
            bundle_info: bundle_info.into(),
            archetype_move_type,
            change_tick,
            world: world.as_unsafe_world_cell(),
        };

        if is_new_created {
            // SAFETY:
            // - we have exclusive ownership of the world and hold no references to command queue or component data
            // - as this goes through `DeferredWorld`, our pointers will not be invalidated
            let mut world = unsafe { inserter.world.into_deferred() };
            // SAFETY: the ARCHETYPE_CREATED event_key corresponds to the ArchetypeCreated event's type
            unsafe {
                world.trigger_raw(
                    ARCHETYPE_CREATED,
                    &mut ArchetypeCreated(new_archetype_id),
                    &mut GlobalTrigger,
                    MaybeLocation::caller(),
                );
            }
        }
        inserter
    }

    // A non-generic prelude to insert used to minimize duplicated monomorphized code.
    // In combination with after_insert, this can reduce compile time of bevy by 10%.
    // We inline in release to avoid a minor perf loss.
    #[cfg_attr(not(debug_assertions), inline(always))]
    unsafe fn before_insert<'a>(
        entity: Entity,
        location: EntityLocation,
        insert_mode: InsertMode,
        caller: MaybeLocation,
        relationship_hook_mode: RelationshipHookMode,
        mut archetype: NonNull<Archetype>,
        archetype_after_insert: &ArchetypeAfterBundleInsert,
        world: &'a UnsafeWorldCell<'w>,
        archetype_move_type: &'a mut ArchetypeMoveType,
        change_tick: Tick,
    ) -> (
        &'a Archetype,
        EntityLocation,
        &'a mut SparseSets,
        &'a mut Table,
    ) {
        // SAFETY: All components in the bundle are guaranteed to exist in the World
        // as they must be initialized before creating the BundleInfo.
        unsafe {
            // SAFETY: Mutable references do not alias and will be dropped after this block
            let mut deferred_world = world.into_deferred();

            if insert_mode == InsertMode::Replace {
                let archetype = archetype.as_ref();
                let new_archetype = match archetype_move_type {
                    ArchetypeMoveType::SameArchetype => archetype,
                    ArchetypeMoveType::NewArchetypeSameTable { new_archetype }
                    | ArchetypeMoveType::NewArchetypeNewTable { new_archetype, .. } => {
                        new_archetype.as_ref()
                    }
                };
                if archetype.has_discard_observer() {
                    // SAFETY: the DISCARD event_key corresponds to the Discard event's type
                    deferred_world.trigger_raw(
                        DISCARD,
                        &mut DiscardEvent { entity },
                        &mut EntityComponentsTrigger {
                            components: archetype_after_insert.existing(),
                            old_archetype: Some(archetype),
                            new_archetype: Some(new_archetype),
                        },
                        caller,
                    );
                }
                deferred_world.trigger_on_discard(
                    archetype,
                    entity,
                    archetype_after_insert.existing().iter().copied(),
                    caller,
                    relationship_hook_mode,
                );
            }
        }

        // SAFETY: Archetype gets borrowed when running the on_discard observers above,
        // so this reference can only be promoted from shared to &mut down here, after they have been ran.
        // The created reference doesn't escape `BundleInserter::insert`
        let archetype = unsafe { archetype.as_mut() };

        match archetype_move_type {
            ArchetypeMoveType::SameArchetype => {
                let (sparse_sets, table) = {
                    // SAFETY: has read+write perms, we do not use the `&mut World` to access archetypes, world ref will be dropped after this block
                    let world = unsafe { world.world_mut() };
                    (
                        &mut world.storages.sparse_sets,
                        &mut world.storages.tables[archetype.table_id()],
                    )
                };

                (&*archetype, location, sparse_sets, table)
            }
            ArchetypeMoveType::NewArchetypeSameTable { new_archetype } => {
                // SAFETY: Points to valid archetype, doesn't conflict with other references, will not outlive `BundleInserter::insert`
                let new_archetype = unsafe { new_archetype.as_mut() };

                let (sparse_sets, table, entities) = {
                    // SAFETY: has read+write perms, we do not use the `&mut World` to access archetypes, world ref will be dropped after this block
                    let world = unsafe { world.world_mut() };
                    (
                        &mut world.storages.sparse_sets,
                        &mut world.storages.tables[new_archetype.table_id()],
                        &mut world.entities,
                    )
                };

                let result = archetype.swap_remove(location.archetype_row);
                if let Some(swapped_entity) = result.swapped_entity {
                    let swapped_location =
                        // SAFETY: If the swap was successful, swapped_entity must be valid.
                        unsafe { entities.get_spawned(swapped_entity).debug_checked_unwrap() };
                    // SAFETY: entity already existed, its archetype_row just changed
                    unsafe {
                        entities.update_existing_location(
                            swapped_entity.index(),
                            Some(EntityLocation {
                                archetype_id: swapped_location.archetype_id,
                                archetype_row: location.archetype_row,
                                table_id: swapped_location.table_id,
                                table_row: swapped_location.table_row,
                            }),
                        )
                    };
                }
                // SAFETY: Component data already existed there
                let new_location = unsafe { new_archetype.allocate(entity, result.table_row) };
                // SAFETY: Entity and therefore index already exists, location was just allocated
                unsafe { entities.update_existing_location(entity.index(), Some(new_location)) };

                (&*new_archetype, new_location, sparse_sets, table)
            }
            ArchetypeMoveType::NewArchetypeNewTable { new_archetype } => {
                // SAFETY: Points to valid archetype, doesn't conflict with other references, will not outlive `BundleInserter::insert`
                let new_archetype = unsafe { new_archetype.as_mut() };

                // SAFETY: Mutable references do not alias and will be dropped after this block
                let (archetypes_ptr, tables, sparse_sets, entities) = {
                    // SAFETY: has read+write perms, we do not use the `&mut World` to access archetypes, world ref will be dropped after this block
                    let world = unsafe { world.world_mut() };
                    let archetype_ptr: *mut Archetype = world.archetypes.archetypes.as_mut_ptr();
                    (
                        archetype_ptr,
                        &mut world.storages.tables,
                        &mut world.storages.sparse_sets,
                        &mut world.entities,
                    )
                };
                let result = archetype.swap_remove(location.archetype_row);
                if let Some(swapped_entity) = result.swapped_entity {
                    let swapped_location =
                        // SAFETY: If the swap was successful, swapped_entity must be valid.
                        unsafe { entities.get_spawned(swapped_entity).debug_checked_unwrap() };
                    // SAFETY: Existing entity has changed archetype row
                    unsafe {
                        entities.update_existing_location(
                            swapped_entity.index(),
                            Some(EntityLocation {
                                archetype_id: swapped_location.archetype_id,
                                archetype_row: location.archetype_row,
                                table_id: swapped_location.table_id,
                                table_row: swapped_location.table_row,
                            }),
                        )
                    };
                }

                // SAFETY:
                // - Since `ArchetypeMoveType` is `NewArchetypeNewTable`,
                //   the entity's old table was determined to be different
                //   from its new table in `Self::new_with_id`.
                // - `location.table_id` is obtained from a valid location
                //   and `new_archetype.table_id()` is obtained from a valid archetype.
                // - `result.table_row` was obtained from a valid archetype.
                // - We will not drop any components.
                // - Valid values will be written to new components by the caller (`Self::insert`).
                let move_result = unsafe {
                    tables.move_row::<true>(
                        location.table_id,
                        new_archetype.table_id(),
                        result.table_row,
                        change_tick,
                    )
                };

                // SAFETY: Table data has been moved to this table/row, sparse set data was already valid
                let new_location = unsafe { new_archetype.allocate(entity, move_result.new_row) };
                // SAFETY: Entity exists & location was just allocated
                unsafe { entities.update_existing_location(entity.index(), Some(new_location)) };

                // If an entity was moved into this entity's table spot, update its table row.
                if let Some(swapped_entity) = move_result.swapped_entity {
                    let swapped_location =
                        // SAFETY: If the swap was successful, swapped_entity must be valid.
                        unsafe { entities.get_spawned(swapped_entity).debug_checked_unwrap() };

                    // SAFETY: Existing entity has changed archetype row
                    unsafe {
                        entities.update_existing_location(
                            swapped_entity.index(),
                            Some(EntityLocation {
                                archetype_id: swapped_location.archetype_id,
                                archetype_row: swapped_location.archetype_row,
                                table_id: swapped_location.table_id,
                                table_row: result.table_row,
                            }),
                        )
                    };

                    if archetype.id() == swapped_location.archetype_id {
                        archetype
                            .set_entity_table_row(swapped_location.archetype_row, result.table_row);
                    } else if new_archetype.id() == swapped_location.archetype_id {
                        new_archetype
                            .set_entity_table_row(swapped_location.archetype_row, result.table_row);
                    } else {
                        // SAFETY: the only two borrowed archetypes are above and we just did collision checks
                        unsafe {
                            (*archetypes_ptr.add(swapped_location.archetype_id.index()))
                                .set_entity_table_row(
                                    swapped_location.archetype_row,
                                    result.table_row,
                                );
                        }
                    }
                }

                (
                    &*new_archetype,
                    new_location,
                    sparse_sets,
                    move_result.new_table,
                )
            }
        }
    }

    /// # Safety
    /// - `entity` must currently exist in the source archetype for this inserter.
    /// - `location` must be `entity`'s location in the archetype.
    /// - `T` must match this [`BundleInserter`] type used to create
    /// - If `T::Effect: !NoBundleEffect.`, then [`apply_effect`] must be called exactly once on
    ///   `bundle` after this function before returning to user-space safe code.
    /// - The value pointed to by `bundle` must not be accessed for anything other than [`apply_effect`]
    ///   or dropped.
    ///
    /// [`apply_effect`]: crate::bundle::DynamicBundle::apply_effect
    #[inline]
    pub(crate) unsafe fn insert<T: DynamicBundle>(
        &mut self,
        entity: Entity,
        location: EntityLocation,
        bundle: MovingPtr<'_, T>,
        insert_mode: InsertMode,
        caller: MaybeLocation,
        relationship_hook_mode: RelationshipHookMode,
    ) -> EntityLocation {
        // SAFETY: Points to valid data; the reference doesn't escape this function.
        // This points to data in the world, but
        // - We have exclusive ownership of the world
        // - we do not construct any other references into world.archetype.edges
        // - we run hooks, which can immutably access `Archetypes`, but cannot get a reference to edges through that
        let archetype_after_insert = unsafe { self.archetype_after_insert.as_ref() };

        let (new_archetype, new_location) = {
            // Non-generic prelude extracted to improve compile time by minimizing monomorphized code.
            // SAFETY: before_insert is functionally part of this function
            let (new_archetype, new_location, sparse_sets, table) = unsafe {
                Self::before_insert(
                    entity,
                    location,
                    insert_mode,
                    caller,
                    relationship_hook_mode,
                    self.archetype,
                    archetype_after_insert,
                    &self.world,
                    &mut self.archetype_move_type,
                    self.change_tick,
                )
            };

            // SAFETY:
            // - bundle component status determined in `insert_bundle_into_archetype`
            // - `apply_effect` called if needed per precondition
            // - table_row & bundle match
            unsafe {
                self.bundle_info.as_ref().write_components(
                    table,
                    sparse_sets,
                    archetype_after_insert,
                    archetype_after_insert.required_components.iter(),
                    entity,
                    new_location.table_row,
                    self.change_tick,
                    bundle,
                    insert_mode,
                    caller,
                );
            }

            (new_archetype, new_location)
        };

        // SAFETY: We have no outstanding mutable references to world as they were dropped
        let deferred_world = unsafe { self.world.into_deferred() };

        // SAFETY: No existing mutable references to archetypes or world
        let old_archetype = unsafe { self.archetype.as_ref() };

        // Non-generic postlude extracted to improve compile time by minimizing monomorphized code.
        Self::after_insert(
            entity,
            insert_mode,
            caller,
            relationship_hook_mode,
            archetype_after_insert,
            old_archetype,
            new_archetype,
            deferred_world,
        );

        new_location
    }

    // A non-generic postlude to insert used to minimize duplicated monomorphized code.
    // In combination with before_insert, this can reduce compile time of bevy by 10%.
    // We inline in release to avoid a minor perf loss.
    #[cfg_attr(not(debug_assertions), inline(always))]
    fn after_insert(
        entity: Entity,
        insert_mode: InsertMode,
        caller: MaybeLocation,
        relationship_hook_mode: RelationshipHookMode,
        archetype_after_insert: &ArchetypeAfterBundleInsert,
        old_archetype: &Archetype,
        new_archetype: &Archetype,
        mut deferred_world: crate::world::DeferredWorld<'_>,
    ) {
        // SAFETY: All components in the bundle are guaranteed to exist in the World
        // as they must be initialized before creating the BundleInfo.
        unsafe {
            deferred_world.trigger_on_add(
                new_archetype,
                entity,
                archetype_after_insert.added().iter().copied(),
                caller,
            );
            if new_archetype.has_add_observer() {
                // SAFETY: the ADD event_key corresponds to the Add event's type
                deferred_world.trigger_raw(
                    ADD,
                    &mut AddEvent { entity },
                    &mut EntityComponentsTrigger {
                        components: archetype_after_insert.added(),
                        old_archetype: Some(old_archetype),
                        new_archetype: Some(new_archetype),
                    },
                    caller,
                );
            }
            match insert_mode {
                InsertMode::Replace => {
                    // Insert triggers for both new and existing components if we're replacing them.
                    deferred_world.trigger_on_insert(
                        new_archetype,
                        entity,
                        archetype_after_insert.inserted().iter().copied(),
                        caller,
                        relationship_hook_mode,
                    );
                    if new_archetype.has_insert_observer() {
                        // SAFETY: the INSERT event_key corresponds to the Insert event's type
                        deferred_world.trigger_raw(
                            INSERT,
                            &mut InsertEvent { entity },
                            &mut EntityComponentsTrigger {
                                components: archetype_after_insert.inserted(),
                                old_archetype: Some(old_archetype),
                                new_archetype: Some(new_archetype),
                            },
                            caller,
                        );
                    }
                }
                InsertMode::Keep => {
                    // Insert triggers only for new components if we're not replacing them (since
                    // nothing is actually inserted).
                    deferred_world.trigger_on_insert(
                        new_archetype,
                        entity,
                        archetype_after_insert.added().iter().copied(),
                        caller,
                        relationship_hook_mode,
                    );
                    if new_archetype.has_insert_observer() {
                        // SAFETY: the INSERT event_key corresponds to the Insert event's type
                        deferred_world.trigger_raw(
                            INSERT,
                            &mut InsertEvent { entity },
                            &mut EntityComponentsTrigger {
                                components: archetype_after_insert.added(),
                                old_archetype: Some(old_archetype),
                                new_archetype: Some(new_archetype),
                            },
                            caller,
                        );
                    }
                }
            }
        }
    }

    #[inline]
    pub(crate) fn entities(&mut self) -> &mut Entities {
        // SAFETY: No outstanding references to self.world, changes to entities cannot invalidate our internal pointers
        unsafe { &mut self.world.world_mut().entities }
    }
}

impl BundleInfo {
    /// Inserts a bundle into the given archetype and returns the resulting archetype and whether a new archetype was created.
    /// This could be the same [`ArchetypeId`], in the event that inserting the given bundle
    /// does not result in an [`Archetype`] change.
    ///
    /// Results are cached in the [`Archetype`] graph to avoid redundant work.
    ///
    /// # Safety
    /// `components` must be the same components as passed in [`Self::new`]
    pub(crate) unsafe fn insert_bundle_into_archetype(
        &self,
        archetypes: &mut Archetypes,
        storages: &mut Storages,
        components: &Components,
        observers: &Observers,
        archetype_id: ArchetypeId,
    ) -> (ArchetypeId, bool) {
        if let Some(archetype_after_insert_id) = archetypes[archetype_id]
            .edges()
            .get_archetype_after_bundle_insert(self.id)
        {
            return (archetype_after_insert_id, false);
        }
        let mut new_table_components = Vec::new();
        let mut new_sparse_set_components = Vec::new();
        let mut bundle_status = Vec::with_capacity(self.explicit_components_len());
        let mut added_required_components = Vec::new();
        let mut added = Vec::new();
        let mut existing = Vec::new();

        let current_archetype = &mut archetypes[archetype_id];
        for component_id in self.iter_explicit_components() {
            if current_archetype.contains(component_id) {
                bundle_status.push(ComponentStatus::Existing);
                existing.push(component_id);
            } else {
                bundle_status.push(ComponentStatus::Added);
                added.push(component_id);
                // SAFETY: component_id exists
                let component_info = unsafe { components.get_info_unchecked(component_id) };
                match component_info.storage_type() {
                    StorageType::Table => new_table_components.push(component_id),
                    StorageType::SparseSet => new_sparse_set_components.push(component_id),
                }
            }
        }

        for (index, component_id) in self.iter_required_components().enumerate() {
            if !current_archetype.contains(component_id) {
                added_required_components.push(self.required_component_constructors[index].clone());
                added.push(component_id);
                // SAFETY: component_id exists
                let component_info = unsafe { components.get_info_unchecked(component_id) };
                match component_info.storage_type() {
                    StorageType::Table => {
                        new_table_components.push(component_id);
                    }
                    StorageType::SparseSet => {
                        new_sparse_set_components.push(component_id);
                    }
                }
            }
        }

        if new_table_components.is_empty() && new_sparse_set_components.is_empty() {
            let edges = current_archetype.edges_mut();
            // The archetype does not change when we insert this bundle.
            edges.cache_archetype_after_bundle_insert(
                self.id,
                archetype_id,
                bundle_status,
                added_required_components,
                added,
                existing,
            );
            (archetype_id, false)
        } else {
            let table_id;
            let table_components;
            let sparse_set_components;
            // The archetype changes when we insert this bundle. Prepare the new archetype and storages.
            {
                let current_archetype = &archetypes[archetype_id];
                table_components = if new_table_components.is_empty() {
                    // If there are no new table components, we can keep using this table.
                    table_id = current_archetype.table_id();
                    current_archetype.table_components().collect()
                } else {
                    new_table_components.extend(current_archetype.table_components());
                    // Sort to ignore order while hashing.
                    new_table_components.sort_unstable();
                    // SAFETY: all component ids in `new_table_components` exist
                    table_id = unsafe {
                        storages
                            .tables
                            .get_id_or_insert(&new_table_components, components)
                    };

                    new_table_components
                };

                sparse_set_components = if new_sparse_set_components.is_empty() {
                    current_archetype.sparse_set_components().collect()
                } else {
                    new_sparse_set_components.extend(current_archetype.sparse_set_components());
                    // Sort to ignore order while hashing.
                    new_sparse_set_components.sort_unstable();
                    new_sparse_set_components
                };
            };
            // SAFETY: ids in self must be valid
            let (new_archetype_id, is_new_created) = unsafe {
                archetypes.get_id_or_insert(
                    components,
                    observers,
                    table_id,
                    table_components,
                    sparse_set_components,
                )
            };

            // Add an edge from the old archetype to the new archetype.
            archetypes[archetype_id]
                .edges_mut()
                .cache_archetype_after_bundle_insert(
                    self.id,
                    new_archetype_id,
                    bundle_status,
                    added_required_components,
                    added,
                    existing,
                );
            (new_archetype_id, is_new_created)
        }
    }
}