Skip to content

Commit 85a38ea

Browse files
committed
Added Enum::variant_index
1 parent eeab710 commit 85a38ea

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.variants() {
284+
for (variant_index, variant) in reflect_enum.variants().iter().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
@@ -77,6 +77,7 @@ impl From<()> for DynamicVariant {
7777
pub struct DynamicEnum {
7878
name: String,
7979
variant_name: String,
80+
variant_index: usize,
8081
variant: DynamicVariant,
8182
}
8283

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

145+
/// Set the current enum variant represented by this struct along with its variant index.
146+
pub fn set_variant_with_index<I: Into<String>, V: Into<DynamicVariant>>(
147+
&mut self,
148+
variant_index: usize,
149+
name: I,
150+
variant: V,
151+
) {
152+
self.variant_index = variant_index;
153+
self.variant_name = name.into();
154+
self.variant = variant.into();
155+
}
156+
120157
/// Create a [`DynamicEnum`] from an existing one.
121158
///
122159
/// This is functionally the same as [`DynamicEnum::from_ref`] except it takes an owned value.
@@ -129,8 +166,9 @@ impl DynamicEnum {
129166
/// This is functionally the same as [`DynamicEnum::from`] except it takes a reference.
130167
pub fn from_ref<TEnum: Enum>(value: &TEnum) -> Self {
131168
match value.variant_type() {
132-
VariantType::Unit => DynamicEnum::new(
169+
VariantType::Unit => DynamicEnum::new_with_index(
133170
value.type_name(),
171+
value.variant_index(),
134172
value.variant_name(),
135173
DynamicVariant::Unit,
136174
),
@@ -139,8 +177,9 @@ impl DynamicEnum {
139177
for field in value.iter_fields() {
140178
data.insert_boxed(field.value().clone_value());
141179
}
142-
DynamicEnum::new(
180+
DynamicEnum::new_with_index(
143181
value.type_name(),
182+
value.variant_index(),
144183
value.variant_name(),
145184
DynamicVariant::Tuple(data),
146185
)
@@ -151,8 +190,9 @@ impl DynamicEnum {
151190
let name = field.name().unwrap();
152191
data.insert_boxed(name, field.value().clone_value());
153192
}
154-
DynamicEnum::new(
193+
DynamicEnum::new_with_index(
155194
value.type_name(),
195+
value.variant_index(),
156196
value.variant_name(),
157197
DynamicVariant::Struct(data),
158198
)
@@ -226,6 +266,10 @@ impl Enum for DynamicEnum {
226266
&self.variant_name
227267
}
228268

269+
fn variant_index(&self) -> usize {
270+
self.variant_index
271+
}
272+
229273
fn variant_type(&self) -> VariantType {
230274
match &self.variant {
231275
DynamicVariant::Unit => VariantType::Unit,
@@ -237,6 +281,7 @@ impl Enum for DynamicEnum {
237281
fn clone_dynamic(&self) -> DynamicEnum {
238282
Self {
239283
name: self.name.clone(),
284+
variant_index: self.variant_index,
240285
variant_name: self.variant_name.clone(),
241286
variant: self.variant.clone(),
242287
}

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
@@ -303,22 +303,15 @@ impl<'a> Serialize for EnumSerializer<'a> {
303303
};
304304

305305
let enum_name = enum_info.name();
306+
let variant_index = self.enum_value.variant_index() as u32;
306307
let variant_info = enum_info
307-
.variant(self.enum_value.variant_name())
308+
.variant_at(variant_index as usize)
308309
.ok_or_else(|| {
309310
Error::custom(format_args!(
310-
"variant `{}` does not exist",
311-
self.enum_value.variant_name()
311+
"variant at index `{}` does not exist",
312+
variant_index
312313
))
313314
})?;
314-
let variant_index = enum_info
315-
.index_of(self.enum_value.variant_name())
316-
.ok_or_else(|| {
317-
Error::custom(format_args!(
318-
"variant `{}` does not exist",
319-
self.enum_value.variant_name()
320-
))
321-
})? as u32;
322315
let variant_name = variant_info.name();
323316
let variant_type = self.enum_value.variant_type();
324317
let field_len = self.enum_value.field_len();

0 commit comments

Comments
 (0)