From f005a96dd4a9aa019fecd5ac874cbd2e962e25f2 Mon Sep 17 00:00:00 2001
From: Benjamin Brienen <benjamin.brienen@outlook.com>
Date: Mon, 28 Oct 2024 23:29:05 +0100
Subject: [PATCH] `ReflectBundle::remove` improvement (#16139)

# Objective

Fixes #15676

## Solution

`remove` returns the removed item

Add `take`

## Testing

None yet

## Migration Guide

If you don't need the returned value from `remove`, discard it.
---
 crates/bevy_ecs/src/reflect/bundle.rs | 28 +++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/crates/bevy_ecs/src/reflect/bundle.rs b/crates/bevy_ecs/src/reflect/bundle.rs
index 771af49319700..d5d52c734964c 100644
--- a/crates/bevy_ecs/src/reflect/bundle.rs
+++ b/crates/bevy_ecs/src/reflect/bundle.rs
@@ -28,14 +28,16 @@ pub struct ReflectBundle(ReflectBundleFns);
 /// The also [`super::component::ReflectComponentFns`].
 #[derive(Clone)]
 pub struct ReflectBundleFns {
-    /// Function pointer implementing [`ReflectBundle::insert()`].
+    /// Function pointer implementing [`ReflectBundle::insert`].
     pub insert: fn(&mut EntityWorldMut, &dyn PartialReflect, &TypeRegistry),
-    /// Function pointer implementing [`ReflectBundle::apply()`].
+    /// Function pointer implementing [`ReflectBundle::apply`].
     pub apply: fn(EntityMut, &dyn PartialReflect, &TypeRegistry),
-    /// Function pointer implementing [`ReflectBundle::apply_or_insert()`].
+    /// Function pointer implementing [`ReflectBundle::apply_or_insert`].
     pub apply_or_insert: fn(&mut EntityWorldMut, &dyn PartialReflect, &TypeRegistry),
-    /// Function pointer implementing [`ReflectBundle::remove()`].
+    /// Function pointer implementing [`ReflectBundle::remove`].
     pub remove: fn(&mut EntityWorldMut),
+    /// Function pointer implementing [`ReflectBundle::take`].
+    pub take: fn(&mut EntityWorldMut) -> Option<Box<dyn Reflect>>,
 }
 
 impl ReflectBundleFns {
@@ -85,8 +87,17 @@ impl ReflectBundle {
     }
 
     /// Removes this [`Bundle`] type from the entity. Does nothing if it doesn't exist.
-    pub fn remove(&self, entity: &mut EntityWorldMut) {
+    pub fn remove(&self, entity: &mut EntityWorldMut) -> &ReflectBundle {
         (self.0.remove)(entity);
+        self
+    }
+
+    /// Removes all components in the [`Bundle`] from the entity and returns their previous values.
+    ///
+    /// **Note:** If the entity does not have every component in the bundle, this method will not remove any of them.
+    #[must_use]
+    pub fn take(&self, entity: &mut EntityWorldMut) -> Option<Box<dyn Reflect>> {
+        (self.0.take)(entity)
     }
 
     /// Create a custom implementation of [`ReflectBundle`].
@@ -168,7 +179,7 @@ impl<B: Bundle + Reflect + TypePath> FromType<B> for ReflectBundle {
                             .iter_fields()
                             .for_each(|field| apply_or_insert_field(entity, field, registry)),
                         _ => panic!(
-                            "expected bundle `{}` to be named struct or tuple",
+                            "expected bundle `{}` to be a named struct or tuple",
                             // FIXME: once we have unique reflect, use `TypePath`.
                             core::any::type_name::<B>(),
                         ),
@@ -178,6 +189,11 @@ impl<B: Bundle + Reflect + TypePath> FromType<B> for ReflectBundle {
             remove: |entity| {
                 entity.remove::<B>();
             },
+            take: |entity| {
+                entity
+                    .take::<B>()
+                    .map(|bundle| Box::new(bundle).into_reflect())
+            },
         })
     }
 }