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
@@ -300,6 +306,26 @@ impl Archetype {
         }
     }
 
+    /// Swap the entity at `index` with the last one, then decrement `len`
+    ///
+    /// Returns the ID of the entity moved into `index`, if any
+    pub(crate) unsafe fn take(&mut self, index: u32) -> Option<u32> {
+        let last = self.len - 1;
+        if index == last {
+            self.len = last;
+            return None;
+        }
+
+        for (ty, data) in self.types.iter().zip(&*self.data) {
+            let src = data.storage.as_ptr().add(index as usize * ty.layout.size());
+            let moved = data.storage.as_ptr().add(last as usize * ty.layout.size());
+            ptr::swap_nonoverlapping(src, moved, ty.layout.size());
+        }
+        self.len = last;
+        self.entities[index as usize] = self.entities[last as usize];
+        Some(self.entities[last as usize])
+    }
+
     /// Returns the ID of the entity moved into `index`, if any
     pub(crate) unsafe fn remove(&mut self, index: u32, drop: bool) -> Option<u32> {
         let last = self.len - 1;

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,52 @@
+use alloc::vec::Vec;
+use core::mem;
+
+use crate::{Archetype, DynamicBundle, TypeInfo};
+
+/// An entity removed from a `World`
+pub struct TakenEntity<'a> {
+    archetype: &'a Archetype,
+}
+
+impl<'a> TakenEntity<'a> {
+    /// # Safety
+    /// `archetype` must have a valid entity stored past its end
+    pub(crate) unsafe fn new(archetype: &'a Archetype) -> Self {
+        Self { archetype }
+    }
+}
+
+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(self, mut f: impl FnMut(*mut u8, TypeInfo)) {
+        let archetype = self.archetype;
+        mem::forget(self); // Suppress Drop
+        for &ty in archetype.types() {
+            let ptr = archetype
+                .get_dynamic(ty.id(), ty.layout().size(), archetype.len())
+                .unwrap();
+            f(ptr.as_ptr(), ty)
+        }
+    }
+}
+
+impl Drop for TakenEntity<'_> {
+    fn drop(&mut self) {
+        for &ty in self.archetype.types() {
+            unsafe {
+                let ptr = self
+                    .archetype
+                    .get_dynamic(ty.id(), ty.layout().size(), self.archetype.len())
+                    .unwrap();
+                ty.drop(ptr.as_ptr());
+            }
+        }
+    }
+}

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,18 @@ 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> {
+        let loc = self.entities.free(entity)?;
+        let archetype = &mut self.archetypes.archetypes[loc.archetype as usize];
+        if let Some(moved) = unsafe { archetype.take(loc.index) } {
+            self.entities.meta[moved as usize].location.index = loc.index;
+        }
+        unsafe { Ok(TakenEntity::new(archetype)) }
+    }
+
     /// 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));
+}