Add a drain method to all reflect container traits

crates/bevy_reflect/bevy_reflect_derive/src/impls/enums.rs

@@ -3,7 +3,7 @@ use crate::enum_utility::{get_variant_constructors, EnumVariantConstructors};
 use crate::impls::impl_typed;
 use proc_macro::TokenStream;
 use proc_macro2::{Ident, Span};
-use quote::quote;
+use quote::{format_ident, quote};
 
 pub(crate) fn impl_enum(reflect_enum: &ReflectEnum) -> TokenStream {
     let bevy_reflect_path = reflect_enum.meta().bevy_reflect_path();
@@ -22,6 +22,7 @@ pub(crate) fn impl_enum(reflect_enum: &ReflectEnum) -> TokenStream {
         enum_field_len,
         enum_variant_name,
         enum_variant_type,
+        enum_variant_drain,
     } = generate_impls(reflect_enum, &ref_index, &ref_name);
 
     let EnumVariantConstructors {
@@ -120,6 +121,13 @@ pub(crate) fn impl_enum(reflect_enum: &ReflectEnum) -> TokenStream {
                 #bevy_reflect_path::VariantFieldIter::new(self)
             }
 
+            fn drain(self: Box<Self>) -> Vec<Box<dyn #bevy_reflect_path::Reflect>> {
+                match *self {
+                    #(#enum_variant_drain,)*
+                    _ => Vec::new()
+                }
+            }
+
             #[inline]
             fn field_len(&self) -> usize {
                  match self {
@@ -255,6 +263,7 @@ struct EnumImpls {
     enum_field_len: Vec<proc_macro2::TokenStream>,
     enum_variant_name: Vec<proc_macro2::TokenStream>,
     enum_variant_type: Vec<proc_macro2::TokenStream>,
+    enum_variant_drain: Vec<proc_macro2::TokenStream>,
 }
 
 fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Ident) -> EnumImpls {
@@ -268,6 +277,7 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
     let mut enum_field_len = Vec::new();
     let mut enum_variant_name = Vec::new();
     let mut enum_variant_type = Vec::new();
+    let mut enum_variant_drain = Vec::new();
 
     for variant in reflect_enum.active_variants() {
         let ident = &variant.data.ident;
@@ -323,6 +333,19 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
                 });
                 let arguments = quote!(#name, &[ #(#argument),* ]);
                 add_fields_branch("Tuple", "TupleVariantInfo", arguments, field_len);
+
+                let field_idents = fields
+                    .iter()
+                    .filter(|field| !field.attrs.ignore)
+                    .map(|field| syn::Index::from(field.index));
+                let field_names = fields
+                    .iter()
+                    .filter(|field| !field.attrs.ignore)
+                    .map(|field| format_ident!("_{}", field.index))
+                    .collect::<Vec<_>>();
+                enum_variant_drain.push(quote! {
+                   #unit{ #(#field_idents : #field_names,)* .. } => vec![#(Box::new(#field_names) as Box<dyn #bevy_reflect_path::Reflect>),*]
+                });
             }
             EnumVariantFields::Named(fields) => {
                 let (field_len, argument) = for_fields(fields, |reflect_idx, _, field| {
@@ -346,6 +369,15 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
                 });
                 let arguments = quote!(#name, &[ #(#argument),* ]);
                 add_fields_branch("Struct", "StructVariantInfo", arguments, field_len);
+
+                let field_idents = fields
+                    .iter()
+                    .filter(|field| !field.attrs.ignore)
+                    .map(|field| field.data.ident.as_ref().unwrap())
+                    .collect::<Vec<_>>();
+                enum_variant_drain.push(quote! {
+                   #unit{ #(#field_idents,)* .. } => vec![#(Box::new(#field_idents) as Box<dyn #bevy_reflect_path::Reflect>),*]
+                });
             }
         };
     }
@@ -359,5 +391,6 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
         enum_field_len,
         enum_variant_name,
         enum_variant_type,
+        enum_variant_drain,
     }
 }

crates/bevy_reflect/bevy_reflect_derive/src/impls/structs.rs

@@ -117,6 +117,12 @@ pub(crate) fn impl_struct(reflect_struct: &ReflectStruct) -> TokenStream {
                 #bevy_reflect_path::FieldIter::new(self)
             }
 
+            fn drain(self: Box<Self>) -> Vec<Box<dyn #bevy_reflect_path::Reflect>> {
+                vec![
+                    #(Box::new(self.#field_idents),)*
+                ]
+            }
+
             fn clone_dynamic(&self) -> #bevy_reflect_path::DynamicStruct {
                 let mut dynamic = #bevy_reflect_path::DynamicStruct::default();
                 dynamic.set_name(self.type_name().to_string());

crates/bevy_reflect/bevy_reflect_derive/src/impls/tuple_structs.rs

@@ -79,6 +79,12 @@ pub(crate) fn impl_tuple_struct(reflect_struct: &ReflectStruct) -> TokenStream {
                 #bevy_reflect_path::TupleStructFieldIter::new(self)
             }
 
+            fn drain(self: Box<Self>) -> Vec<Box<dyn #bevy_reflect_path::Reflect>> {
+                vec![
+                    #(Box::new(self.#field_idents),)*
+                ]
+            }
+
             fn clone_dynamic(&self) -> #bevy_reflect_path::DynamicTupleStruct {
                 let mut dynamic = #bevy_reflect_path::DynamicTupleStruct::default();
                 dynamic.set_name(self.type_name().to_string());

crates/bevy_reflect/src/array.rs

@@ -30,6 +30,8 @@ pub trait Array: Reflect {
     }
     /// Returns an iterator over the collection.
     fn iter(&self) -> ArrayIter;
+    /// Drain the elements of this array to get a vector of owned values.
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>>;
 
     fn clone_dynamic(&self) -> DynamicArray {
         DynamicArray {
@@ -246,6 +248,11 @@ impl Array for DynamicArray {
         }
     }
 
+    #[inline]
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.values.into_vec()
+    }
+
     #[inline]
     fn clone_dynamic(&self) -> DynamicArray {
         DynamicArray {

crates/bevy_reflect/src/enums/dynamic_enum.rs

@@ -213,6 +213,14 @@ impl Enum for DynamicEnum {
         VariantFieldIter::new(self)
     }
 
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        match self.variant {
+            DynamicVariant::Unit => Vec::new(),
+            DynamicVariant::Tuple(tuple_variant) => Box::new(tuple_variant).drain(),
+            DynamicVariant::Struct(struct_variant) => Box::new(struct_variant).drain(),
+        }
+    }
+
     fn field_len(&self) -> usize {
         match &self.variant {
             DynamicVariant::Unit => 0,

crates/bevy_reflect/src/enums/enum_trait.rs

@@ -110,6 +110,8 @@ pub trait Enum: Reflect {
     fn name_at(&self, index: usize) -> Option<&str>;
     /// Returns an iterator over the values of the current variant's fields.
     fn iter_fields(&self) -> VariantFieldIter;
+    /// Drain the fields of this enum to get a vector of owned values.
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>>;
     /// Returns the number of fields in the current variant.
     fn field_len(&self) -> usize;
     /// The name of the current variant.

crates/bevy_reflect/src/impls/smallvec.rs

@@ -36,6 +36,12 @@ where
             index: 0,
         }
     }
+
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.into_iter()
+            .map(|value| Box::new(value) as Box<dyn Reflect>)
+            .collect()
+    }
 }
 
 impl<T: smallvec::Array + Send + Sync + 'static> List for SmallVec<T>

crates/bevy_reflect/src/impls/std.rs

@@ -113,6 +113,13 @@ impl<T: FromReflect> Array for Vec<T> {
             index: 0,
         }
     }
+
+    #[inline]
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.into_iter()
+            .map(|value| Box::new(value) as Box<dyn Reflect>)
+            .collect()
+    }
 }
 
 impl<T: FromReflect> List for Vec<T> {
@@ -247,6 +254,17 @@ impl<K: FromReflect + Eq + Hash, V: FromReflect> Map for HashMap<K, V> {
         }
     }
 
+    fn drain(self: Box<Self>) -> Vec<(Box<dyn Reflect>, Box<dyn Reflect>)> {
+        self.into_iter()
+            .map(|(key, value)| {
+                (
+                    Box::new(key) as Box<dyn Reflect>,
+                    Box::new(value) as Box<dyn Reflect>,
+                )
+            })
+            .collect()
+    }
+
     fn clone_dynamic(&self) -> DynamicMap {
         let mut dynamic_map = DynamicMap::default();
         dynamic_map.set_name(self.type_name().to_string());
@@ -395,6 +413,13 @@ impl<T: Reflect, const N: usize> Array for [T; N] {
             index: 0,
         }
     }
+
+    #[inline]
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.into_iter()
+            .map(|value| Box::new(value) as Box<dyn Reflect>)
+            .collect()
+    }
 }
 
 impl<T: Reflect, const N: usize> Reflect for [T; N] {
@@ -627,6 +652,13 @@ impl<T: FromReflect> Enum for Option<T> {
         VariantFieldIter::new(self)
     }
 
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        match *self {
+            Some(value) => vec![Box::new(value) as Box<dyn Reflect>],
+            None => Vec::new(),
+        }
+    }
+
     #[inline]
     fn field_len(&self) -> usize {
         match self {

crates/bevy_reflect/src/lib.rs

@@ -519,6 +519,70 @@ mod tests {
         assert_eq!(foo, *foo2.downcast::<Foo>().unwrap());
     }
 
+    #[test]
+    fn should_drain_fields() {
+        #[derive(Reflect, PartialEq, Debug)]
+        #[reflect(PartialEq, Debug)]
+        struct Foo(usize);
+
+        #[derive(Reflect, PartialEq, Debug)]
+        #[reflect(PartialEq, Debug)]
+        struct Bar {
+            a: i32,
+            b: f32,
+        }
+
+        #[derive(Reflect, PartialEq, Debug)]
+        #[reflect(PartialEq, Debug)]
+        enum Baz {
+            Unit,
+            Tuple(i32, f32),
+            Struct { value: i32 },
+        }
+
+        let array_value: Box<dyn Array> = Box::new([123_i32, 321_i32]);
+        let fields = array_value.drain();
+        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 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());
+
+        let tuple_value: Box<dyn Tuple> = Box::new((123_i32, Foo(1337)));
+        let fields = tuple_value.drain();
+        assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
+        assert!(fields[1].reflect_partial_eq(&Foo(1337)).unwrap_or_default());
+
+        let 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());
+
+        let tuple_struct_value: Box<dyn TupleStruct> = Box::new(Foo(123));
+        let fields = tuple_struct_value.drain();
+        assert!(fields[0].reflect_partial_eq(&123usize).unwrap_or_default());
+
+        let struct_value: Box<dyn Struct> = Box::new(Bar { a: 123, b: 1.23 });
+        let fields = struct_value.drain();
+        assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
+        assert!(fields[1].reflect_partial_eq(&1.23_f32).unwrap_or_default());
+
+        let unit_variant_value: Box<dyn Enum> = Box::new(Baz::Unit);
+        let fields = unit_variant_value.drain();
+        assert_eq!(0, fields.len());
+
+        let tuple_variant_value: Box<dyn Enum> = Box::new(Baz::Tuple(123, 1.23));
+        let fields = tuple_variant_value.drain();
+        assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
+        assert!(fields[1].reflect_partial_eq(&1.23_f32).unwrap_or_default());
+
+        let struct_variant_value: Box<dyn Enum> = Box::new(Baz::Struct { value: 123 });
+        let fields = struct_variant_value.drain();
+        assert!(fields[0].reflect_partial_eq(&123_i32).unwrap_or_default());
+    }
+
     #[test]
     fn reflect_take() {
         #[derive(Reflect, Debug, PartialEq)]

crates/bevy_reflect/src/list.rs

@@ -134,6 +134,10 @@ impl Array for DynamicList {
         }
     }
 
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.values
+    }
+
     fn clone_dynamic(&self) -> DynamicArray {
         DynamicArray {
             name: self.name.clone(),

crates/bevy_reflect/src/map.rs

@@ -42,6 +42,9 @@ pub trait Map: Reflect {
     /// Returns an iterator over the key-value pairs of the map.
     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 Reflect>, Box<dyn Reflect>)>;
+
     /// Clones the map, producing a [`DynamicMap`].
     fn clone_dynamic(&self) -> DynamicMap;
 
@@ -226,6 +229,10 @@ impl Map for DynamicMap {
             }
         }
     }
+
+    fn drain(self: Box<Self>) -> Vec<(Box<dyn Reflect>, Box<dyn Reflect>)> {
+        self.values
+    }
 }
 
 impl Reflect for DynamicMap {

crates/bevy_reflect/src/struct_trait.rs

@@ -62,6 +62,9 @@ pub trait Struct: Reflect {
     /// Returns an iterator over the values of the struct's fields.
     fn iter_fields(&self) -> FieldIter;
 
+    /// Drain the fields of this struct to get a vector of owned values.
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>>;
+
     /// Clones the struct into a [`DynamicStruct`].
     fn clone_dynamic(&self) -> DynamicStruct;
 }
@@ -328,6 +331,11 @@ impl Struct for DynamicStruct {
         }
     }
 
+    #[inline]
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.fields
+    }
+
     fn clone_dynamic(&self) -> DynamicStruct {
         DynamicStruct {
             name: self.name.clone(),

crates/bevy_reflect/src/tuple.rs

@@ -40,6 +40,9 @@ pub trait Tuple: Reflect {
     /// Returns an iterator over the values of the tuple's fields.
     fn iter_fields(&self) -> TupleFieldIter;
 
+    /// Drain the fields of this tuple to get a vector of owned values.
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>>;
+
     /// Clones the struct into a [`DynamicTuple`].
     fn clone_dynamic(&self) -> DynamicTuple;
 }
@@ -253,6 +256,11 @@ impl Tuple for DynamicTuple {
         }
     }
 
+    #[inline]
+    fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+        self.fields
+    }
+
     #[inline]
     fn clone_dynamic(&self) -> DynamicTuple {
         DynamicTuple {
@@ -451,6 +459,13 @@ macro_rules! impl_reflect_tuple {
                 }
             }
 
+            #[inline]
+            fn drain(self: Box<Self>) -> Vec<Box<dyn Reflect>> {
+                vec![
+                    $(Box::new(self.$index),)*
+                ]
+            }
+
             #[inline]
             fn clone_dynamic(&self) -> DynamicTuple {
                 let mut dyn_tuple = DynamicTuple {