Introduce method for moving entities between worlds

CHANGELOG.md

@@ -1,3 +1,8 @@
+# Unreleased
+
+### Added
+- `World::take` for moving complete entities between worlds
+
 # 0.7.5
 
 ### Changed

src/archetype.rs

@@ -26,6 +26,7 @@ use crate::{Access, Component, Query};
 /// [`World`](crate::World).
 pub struct Archetype {
     types: Vec<TypeInfo>,
+    type_ids: Box<[TypeId]>,
     index: OrderedTypeIdMap<usize>,
     len: u32,
     entities: Box<[u32]>,
@@ -58,6 +59,7 @@ impl Archetype {
         let component_count = types.len();
         Self {
             index: OrderedTypeIdMap::new(types.iter().enumerate().map(|(i, ty)| (ty.id, i))),
+            type_ids: types.iter().map(|ty| ty.id()).collect(),
             types,
             entities: Box::new([]),
             len: 0,
@@ -178,6 +180,10 @@ impl Archetype {
         &self.types
     }
 
+    pub(crate) fn type_ids(&self) -> &[TypeId] {
+        &self.type_ids
+    }
+
     /// Enumerate the types of the components of entities stored in this archetype.
     ///
     /// Convenient for dispatching logic which needs to be performed on sets of type ids.  For

src/lib.rs

@@ -69,6 +69,7 @@ mod query;
 mod query_one;
 #[cfg(any(feature = "row-serialize", feature = "column-serialize"))]
 pub mod serialize;
+mod take;
 mod world;
 
 pub use archetype::{Archetype, ArchetypeColumn};
@@ -85,6 +86,7 @@ pub use query::{
     With, Without,
 };
 pub use query_one::QueryOne;
+pub use take::TakenEntity;
 pub use world::{
     ArchetypesGeneration, Component, ComponentError, Iter, QueryOneError, SpawnBatchIter,
     SpawnColumnBatchIter, World,

src/take.rs

@@ -0,0 +1,64 @@
+use alloc::vec::Vec;
+
+use crate::{entities::Entities, Archetype, DynamicBundle, Entity, TypeInfo};
+
+/// An entity removed from a `World`
+pub struct TakenEntity<'a> {
+    entities: &'a mut Entities,
+    entity: Entity,
+    archetype: &'a mut Archetype,
+    index: u32,
+    drop: bool,
+}
+
+impl<'a> TakenEntity<'a> {
+    /// # Safety
+    /// `index` must be in bounds in `archetype`
+    pub(crate) unsafe fn new(
+        entities: &'a mut Entities,
+        entity: Entity,
+        archetype: &'a mut Archetype,
+        index: u32,
+    ) -> Self {
+        Self {
+            entities,
+            entity,
+            archetype,
+            index,
+            drop: true,
+        }
+    }
+}
+
+unsafe impl<'a> DynamicBundle for TakenEntity<'a> {
+    fn with_ids<T>(&self, f: impl FnOnce(&[core::any::TypeId]) -> T) -> T {
+        f(self.archetype.type_ids())
+    }
+
+    fn type_info(&self) -> Vec<crate::TypeInfo> {
+        self.archetype.types().to_vec()
+    }
+
+    unsafe fn put(mut self, mut f: impl FnMut(*mut u8, TypeInfo)) {
+        // Suppress dropping of moved components
+        self.drop = false;
+        for &ty in self.archetype.types() {
+            let ptr = self
+                .archetype
+                .get_dynamic(ty.id(), ty.layout().size(), self.index)
+                .unwrap();
+            f(ptr.as_ptr(), ty)
+        }
+    }
+}
+
+impl Drop for TakenEntity<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            if let Some(moved) = self.archetype.remove(self.index, self.drop) {
+                self.entities.meta[moved as usize].location.index = self.index;
+            }
+            self.entities.free(self.entity).unwrap();
+        }
+    }
+}

src/world.rs

@@ -26,6 +26,7 @@ use crate::entities::{Entities, EntityMeta, Location, ReserveEntitiesIterator};
 use crate::{
     Bundle, Column, ColumnBatch, ColumnMut, DynamicBundle, Entity, EntityRef, Fetch,
     MissingComponent, NoSuchEntity, Query, QueryBorrow, QueryItem, QueryMut, QueryOne, Ref, RefMut,
+    TakenEntity,
 };
 
 /// An unordered collection of entities, each having any number of distinctly typed components
@@ -302,6 +303,8 @@ impl World {
     }
 
     /// Destroy an entity and all its components
+    ///
+    /// See also [`take`](Self::take).
     pub fn despawn(&mut self, entity: Entity) -> Result<(), NoSuchEntity> {
         self.flush();
         let loc = self.entities.free(entity)?;
@@ -878,6 +881,23 @@ impl World {
         ColumnMut::new(entities, archetypes, PhantomData)
     }
 
+    /// Despawn `entity`, yielding a [`DynamicBundle`] of its components
+    ///
+    /// Useful for moving entities between worlds.
+    pub fn take(&mut self, entity: Entity) -> Result<TakenEntity<'_>, NoSuchEntity> {
+        self.flush();
+        let loc = self.entities.get(entity)?;
+        let archetype = &mut self.archetypes.archetypes[loc.archetype as usize];
+        unsafe {
+            Ok(TakenEntity::new(
+                &mut self.entities,
+                entity,
+                archetype,
+                loc.index,
+            ))
+        }
+    }
+
     /// Returns a distinct value after `archetypes` is changed
     ///
     /// Store the current value after deriving information from [`archetypes`](Self::archetypes),

tests/tests.rs

@@ -820,3 +820,19 @@ fn len() {
     world.clear();
     assert_eq!(world.len(), 0);
 }
+
+#[test]
+fn take() {
+    let mut world_a = World::new();
+    let e = world_a.spawn(("abc".to_string(), 42));
+    let f = world_a.spawn(("def".to_string(), 17));
+    let mut world_b = World::new();
+    let e2 = world_b.spawn(world_a.take(e).unwrap());
+    assert!(!world_a.contains(e));
+    assert_eq!(*world_b.get::<String>(e2).unwrap(), "abc");
+    assert_eq!(*world_b.get::<i32>(e2).unwrap(), 42);
+    assert_eq!(*world_a.get::<String>(f).unwrap(), "def");
+    assert_eq!(*world_a.get::<i32>(f).unwrap(), 17);
+    world_b.take(e2).unwrap();
+    assert!(!world_b.contains(e2));
+}