Skip to content

Commit 9ebd312

Browse files
committed
Added Enum::variant_index
1 parent 88a6117 commit 9ebd312

File tree

5 files changed

+81
-18
lines changed

5 files changed

+81
-18
lines changed

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

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ pub(crate) fn impl_enum(reflect_enum: &ReflectEnum) -> TokenStream {
2121
enum_name_at,
2222
enum_field_len,
2323
enum_variant_name,
24+
enum_variant_index,
2425
enum_variant_type,
2526
} = generate_impls(reflect_enum, &ref_index, &ref_name);
2627

@@ -137,6 +138,14 @@ pub(crate) fn impl_enum(reflect_enum: &ReflectEnum) -> TokenStream {
137138
}
138139
}
139140

141+
#[inline]
142+
fn variant_index(&self) -> usize {
143+
match self {
144+
#(#enum_variant_index,)*
145+
_ => unreachable!(),
146+
}
147+
}
148+
140149
#[inline]
141150
fn variant_type(&self) -> #bevy_reflect_path::VariantType {
142151
match self {
@@ -255,6 +264,7 @@ struct EnumImpls {
255264
enum_name_at: Vec<proc_macro2::TokenStream>,
256265
enum_field_len: Vec<proc_macro2::TokenStream>,
257266
enum_variant_name: Vec<proc_macro2::TokenStream>,
267+
enum_variant_index: Vec<proc_macro2::TokenStream>,
258268
enum_variant_type: Vec<proc_macro2::TokenStream>,
259269
}
260270

@@ -268,13 +278,21 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
268278
let mut enum_name_at = Vec::new();
269279
let mut enum_field_len = Vec::new();
270280
let mut enum_variant_name = Vec::new();
281+
let mut enum_variant_index = Vec::new();
271282
let mut enum_variant_type = Vec::new();
272283

273-
for variant in reflect_enum.active_variants() {
284+
for (variant_index, variant) in reflect_enum.active_variants().enumerate() {
274285
let ident = &variant.data.ident;
275286
let name = ident.to_string();
276287
let unit = reflect_enum.get_unit(ident);
277288

289+
enum_variant_name.push(quote! {
290+
#unit{..} => #name
291+
});
292+
enum_variant_index.push(quote! {
293+
#unit{..} => #variant_index
294+
});
295+
278296
fn for_fields(
279297
fields: &[StructField],
280298
mut generate_for_field: impl FnMut(usize, usize, &StructField) -> proc_macro2::TokenStream,
@@ -302,9 +320,6 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
302320
enum_field_len.push(quote! {
303321
#unit{..} => #field_len
304322
});
305-
enum_variant_name.push(quote! {
306-
#unit{..} => #name
307-
});
308323
enum_variant_type.push(quote! {
309324
#unit{..} => #bevy_reflect_path::VariantType::#variant
310325
});
@@ -359,6 +374,7 @@ fn generate_impls(reflect_enum: &ReflectEnum, ref_index: &Ident, ref_name: &Iden
359374
enum_name_at,
360375
enum_field_len,
361376
enum_variant_name,
377+
enum_variant_index,
362378
enum_variant_type,
363379
}
364380
}

crates/bevy_reflect/src/enums/dynamic_enum.rs

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ impl From<()> for DynamicVariant {
7676
pub struct DynamicEnum {
7777
name: String,
7878
variant_name: String,
79+
variant_index: usize,
7980
variant: DynamicVariant,
8081
}
8182

@@ -95,6 +96,30 @@ impl DynamicEnum {
9596
) -> Self {
9697
Self {
9798
name: name.into(),
99+
variant_index: 0,
100+
variant_name: variant_name.into(),
101+
variant: variant.into(),
102+
}
103+
}
104+
105+
/// Create a new [`DynamicEnum`] with a variant index to represent an enum at runtime.
106+
///
107+
/// # Arguments
108+
///
109+
/// * `name`: The type name of the enum
110+
/// * `variant_index`: The index of the variant to set
111+
/// * `variant_name`: The name of the variant to set
112+
/// * `variant`: The variant data
113+
///
114+
pub fn new_with_index<I: Into<String>, V: Into<DynamicVariant>>(
115+
name: I,
116+
variant_index: usize,
117+
variant_name: I,
118+
variant: V,
119+
) -> Self {
120+
Self {
121+
name: name.into(),
122+
variant_index,
98123
variant_name: variant_name.into(),
99124
variant: variant.into(),
100125
}
@@ -116,6 +141,18 @@ impl DynamicEnum {
116141
self.variant = variant.into();
117142
}
118143

144+
/// Set the current enum variant represented by this struct along with its variant index.
145+
pub fn set_variant_with_index<I: Into<String>, V: Into<DynamicVariant>>(
146+
&mut self,
147+
variant_index: usize,
148+
name: I,
149+
variant: V,
150+
) {
151+
self.variant_index = variant_index;
152+
self.variant_name = name.into();
153+
self.variant = variant.into();
154+
}
155+
119156
/// Create a [`DynamicEnum`] from an existing one.
120157
///
121158
/// This is functionally the same as [`DynamicEnum::from_ref`] except it takes an owned value.
@@ -128,8 +165,9 @@ impl DynamicEnum {
128165
/// This is functionally the same as [`DynamicEnum::from`] except it takes a reference.
129166
pub fn from_ref<TEnum: Enum>(value: &TEnum) -> Self {
130167
match value.variant_type() {
131-
VariantType::Unit => DynamicEnum::new(
168+
VariantType::Unit => DynamicEnum::new_with_index(
132169
value.type_name(),
170+
value.variant_index(),
133171
value.variant_name(),
134172
DynamicVariant::Unit,
135173
),
@@ -138,8 +176,9 @@ impl DynamicEnum {
138176
for field in value.iter_fields() {
139177
data.insert_boxed(field.value().clone_value());
140178
}
141-
DynamicEnum::new(
179+
DynamicEnum::new_with_index(
142180
value.type_name(),
181+
value.variant_index(),
143182
value.variant_name(),
144183
DynamicVariant::Tuple(data),
145184
)
@@ -150,8 +189,9 @@ impl DynamicEnum {
150189
let name = field.name().unwrap();
151190
data.insert_boxed(name, field.value().clone_value());
152191
}
153-
DynamicEnum::new(
192+
DynamicEnum::new_with_index(
154193
value.type_name(),
194+
value.variant_index(),
155195
value.variant_name(),
156196
DynamicVariant::Struct(data),
157197
)
@@ -225,6 +265,10 @@ impl Enum for DynamicEnum {
225265
&self.variant_name
226266
}
227267

268+
fn variant_index(&self) -> usize {
269+
self.variant_index
270+
}
271+
228272
fn variant_type(&self) -> VariantType {
229273
match &self.variant {
230274
DynamicVariant::Unit => VariantType::Unit,
@@ -236,6 +280,7 @@ impl Enum for DynamicEnum {
236280
fn clone_dynamic(&self) -> DynamicEnum {
237281
Self {
238282
name: self.name.clone(),
283+
variant_index: self.variant_index,
239284
variant_name: self.variant_name.clone(),
240285
variant: self.variant.clone(),
241286
}

crates/bevy_reflect/src/enums/enum_trait.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ pub trait Enum: Reflect {
113113
fn field_len(&self) -> usize;
114114
/// The name of the current variant.
115115
fn variant_name(&self) -> &str;
116+
/// The index of the current variant.
117+
fn variant_index(&self) -> usize;
116118
/// The type of the current variant.
117119
fn variant_type(&self) -> VariantType;
118120
// Clones the enum into a [`DynamicEnum`].

crates/bevy_reflect/src/impls/std.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -682,6 +682,13 @@ impl<T: FromReflect> Enum for Option<T> {
682682
}
683683
}
684684

685+
fn variant_index(&self) -> usize {
686+
match self {
687+
None => 0,
688+
Some(..) => 1,
689+
}
690+
}
691+
685692
#[inline]
686693
fn variant_type(&self) -> VariantType {
687694
match self {

crates/bevy_reflect/src/serde/ser.rs

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -272,22 +272,15 @@ impl<'a> Serialize for EnumSerializer<'a> {
272272
};
273273

274274
let enum_name = enum_info.name();
275+
let variant_index = self.enum_value.variant_index() as u32;
275276
let variant_info = enum_info
276-
.variant(self.enum_value.variant_name())
277+
.variant_at(variant_index as usize)
277278
.ok_or_else(|| {
278279
Error::custom(format_args!(
279-
"variant `{}` does not exist",
280-
self.enum_value.variant_name()
280+
"variant at index `{}` does not exist",
281+
variant_index
281282
))
282283
})?;
283-
let variant_index = enum_info
284-
.index_of(self.enum_value.variant_name())
285-
.ok_or_else(|| {
286-
Error::custom(format_args!(
287-
"variant `{}` does not exist",
288-
self.enum_value.variant_name()
289-
))
290-
})? as u32;
291284
let variant_name = variant_info.name();
292285
let variant_type = self.enum_value.variant_type();
293286
let field_len = self.enum_value.field_len();

0 commit comments

Comments
 (0)