Make drain take a mutable borrow instead of Box<Self> for reflect…

…ed `Map`, `List`, and `Set`. Some notable exceptions to this change are `Array` and `Tuple`. These types don't make sense with `drain` taking a mutable borrow since they can't get "smaller".
bevyengine · Sep 24, 2024 · 5cce145 · 5cce145
1 parent 3d0f240
commit 5cce145
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 34 deletions.
diff --git a/crates/bevy_reflect/src/impls/smallvec.rs b/crates/bevy_reflect/src/impls/smallvec.rs
@@ -71,8 +71,8 @@ where
         ListIter::new(self)
     }
 
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-        self.into_iter()
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+        self.drain(..)
             .map(|value| Box::new(value) as Box<dyn PartialReflect>)
             .collect()
     }

diff --git a/crates/bevy_reflect/src/impls/std.rs b/crates/bevy_reflect/src/impls/std.rs
@@ -444,8 +444,8 @@ macro_rules! impl_reflect_for_veclike {
             }
 
             #[inline]
-            fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-                self.into_iter()
+            fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+                self.drain(..)
                     .map(|value| Box::new(value) as Box<dyn PartialReflect>)
                     .collect()
             }
@@ -621,8 +621,8 @@ macro_rules! impl_reflect_for_hashmap {
                 MapIter::new(self)
             }
 
-            fn drain(self: Box<Self>) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
-                self.into_iter()
+            fn drain(&mut self) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
+                self.drain()
                     .map(|(key, value)| {
                         (
                             Box::new(key) as Box<dyn PartialReflect>,
@@ -858,8 +858,8 @@ macro_rules! impl_reflect_for_hashset {
                 Box::new(iter)
             }
 
-            fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-                self.into_iter()
+            fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+                self.drain()
                     .map(|value| Box::new(value) as Box<dyn PartialReflect>)
                     .collect()
             }
@@ -1094,15 +1094,18 @@ where
         MapIter::new(self)
     }
 
-    fn drain(self: Box<Self>) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
-        self.into_iter()
-            .map(|(key, value)| {
-                (
-                    Box::new(key) as Box<dyn PartialReflect>,
-                    Box::new(value) as Box<dyn PartialReflect>,
-                )
-            })
-            .collect()
+    fn drain(&mut self) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
+        // BTreeMap doesn't have a `drain` function. See
+        // https://github.com/rust-lang/rust/issues/81074. So we have to fake one by popping
+        // elements off one at a time.
+        let mut result = Vec::with_capacity(self.len());
+        while let Some((k, v)) = self.pop_first() {
+            result.push((
+                Box::new(k) as Box<dyn PartialReflect>,
+                Box::new(v) as Box<dyn PartialReflect>,
+            ));
+        }
+        result
     }
 
     fn clone_dynamic(&self) -> DynamicMap {
@@ -1688,12 +1691,10 @@ impl<T: FromReflect + MaybeTyped + Clone + TypePath + GetTypeRegistration> List
         ListIter::new(self)
     }
 
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-        // into_owned() is not unnecessary here because it avoids cloning whenever you have a Cow::Owned already
-        #[allow(clippy::unnecessary_to_owned)]
-        self.into_owned()
-            .into_iter()
-            .map(|value| value.clone_value())
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+        self.to_mut()
+            .drain(..)
+            .map(|value| Box::new(value) as Box<dyn PartialReflect>)
             .collect()
     }
 }

diff --git a/crates/bevy_reflect/src/lib.rs b/crates/bevy_reflect/src/lib.rs
@@ -1458,7 +1458,7 @@ mod tests {
         assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
         assert!(fields[1].reflect_partial_eq(&321_i32).unwrap_or_default());
 
-        let list_value: Box<dyn List> = Box::new(vec![123_i32, 321_i32]);
+        let mut list_value: Box<dyn List> = Box::new(vec![123_i32, 321_i32]);
         let fields = list_value.drain();
         assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
         assert!(fields[1].reflect_partial_eq(&321_i32).unwrap_or_default());
@@ -1468,7 +1468,7 @@ mod tests {
         assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
         assert!(fields[1].reflect_partial_eq(&321_i32).unwrap_or_default());
 
-        let map_value: Box<dyn Map> = Box::new(HashMap::from([(123_i32, 321_i32)]));
+        let mut map_value: Box<dyn Map> = Box::new(HashMap::from([(123_i32, 321_i32)]));
         let fields = map_value.drain();
         assert!(fields[0].0.reflect_partial_eq(&123_i32).unwrap_or_default());
         assert!(fields[0].1.reflect_partial_eq(&321_i32).unwrap_or_default());

diff --git a/crates/bevy_reflect/src/list.rs b/crates/bevy_reflect/src/list.rs
@@ -95,7 +95,7 @@ pub trait List: PartialReflect {
     fn iter(&self) -> ListIter;
 
     /// Drain the elements of this list to get a vector of owned values.
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>>;
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>>;
 
     /// Clones the list, producing a [`DynamicList`].
     fn clone_dynamic(&self) -> DynamicList {
@@ -228,8 +228,8 @@ impl List for DynamicList {
         ListIter::new(self)
     }
 
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-        self.values
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+        self.values.drain(..).collect()
     }
 
     fn clone_dynamic(&self) -> DynamicList {

diff --git a/crates/bevy_reflect/src/map.rs b/crates/bevy_reflect/src/map.rs
@@ -72,7 +72,7 @@ pub trait Map: PartialReflect {
     fn iter(&self) -> MapIter;
 
     /// Drain the key-value pairs of this map to get a vector of owned values.
-    fn drain(self: Box<Self>) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)>;
+    fn drain(&mut self) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)>;
 
     /// Clones the map, producing a [`DynamicMap`].
     fn clone_dynamic(&self) -> DynamicMap;
@@ -265,8 +265,8 @@ impl Map for DynamicMap {
         MapIter::new(self)
     }
 
-    fn drain(self: Box<Self>) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
-        self.values
+    fn drain(&mut self) -> Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)> {
+        self.values.drain(..).collect()
     }
 
     fn clone_dynamic(&self) -> DynamicMap {

diff --git a/crates/bevy_reflect/src/set.rs b/crates/bevy_reflect/src/set.rs
@@ -58,7 +58,7 @@ pub trait Set: PartialReflect {
     fn iter(&self) -> Box<dyn Iterator<Item = &dyn PartialReflect> + '_>;
 
     /// Drain the values of this set to get a vector of owned values.
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>>;
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>>;
 
     /// Clones the set, producing a [`DynamicSet`].
     fn clone_dynamic(&self) -> DynamicSet;
@@ -184,8 +184,8 @@ impl Set for DynamicSet {
         Box::new(iter)
     }
 
-    fn drain(self: Box<Self>) -> Vec<Box<dyn PartialReflect>> {
-        self.hash_table.into_iter().collect::<Vec<_>>()
+    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
+        self.hash_table.drain().collect::<Vec<_>>()
     }
 
     fn clone_dynamic(&self) -> DynamicSet {