use crate::generics::impl_generic_info_methods;
use crate::{
attributes::{impl_custom_attribute_methods, CustomAttributes},
type_info::impl_type_methods,
DynamicEnum, Generics, PartialReflect, Type, TypePath, VariantInfo, VariantType,
};
use alloc::{boxed::Box, format, string::String};
use bevy_platform::collections::HashMap;
use bevy_platform::sync::Arc;
use core::slice::Iter;
pub trait Enum: PartialReflect {
fn field(&self, name: &str) -> Option<&dyn PartialReflect>;
fn field_at(&self, index: usize) -> Option<&dyn PartialReflect>;
fn field_mut(&mut self, name: &str) -> Option<&mut dyn PartialReflect>;
fn field_at_mut(&mut self, index: usize) -> Option<&mut dyn PartialReflect>;
fn index_of(&self, name: &str) -> Option<usize>;
fn name_at(&self, index: usize) -> Option<&str>;
fn iter_fields(&self) -> VariantFieldIter<'_>;
fn field_len(&self) -> usize;
fn variant_name(&self) -> &str;
fn variant_index(&self) -> usize;
fn variant_type(&self) -> VariantType;
fn to_dynamic_enum(&self) -> DynamicEnum {
DynamicEnum::from_ref(self)
}
fn is_variant(&self, variant_type: VariantType) -> bool {
self.variant_type() == variant_type
}
fn variant_path(&self) -> String {
format!("{}::{}", self.reflect_type_path(), self.variant_name())
}
fn get_represented_enum_info(&self) -> Option<&'static EnumInfo> {
self.get_represented_type_info()?.as_enum().ok()
}
}
#[derive(Clone, Debug)]
pub struct EnumInfo {
ty: Type,
generics: Generics,
variants: Box<[VariantInfo]>,
variant_names: Box<[&'static str]>,
variant_indices: HashMap<&'static str, usize>,
custom_attributes: Arc<CustomAttributes>,
#[cfg(feature = "reflect_documentation")]
docs: Option<&'static str>,
}
impl EnumInfo {
pub fn new<TEnum: Enum + TypePath>(variants: &[VariantInfo]) -> Self {
let variant_indices = variants
.iter()
.enumerate()
.map(|(index, variant)| (variant.name(), index))
.collect::<HashMap<_, _>>();
let variant_names = variants.iter().map(VariantInfo::name).collect();
Self {
ty: Type::of::<TEnum>(),
generics: Generics::new(),
variants: variants.to_vec().into_boxed_slice(),
variant_names,
variant_indices,
custom_attributes: Arc::new(CustomAttributes::default()),
#[cfg(feature = "reflect_documentation")]
docs: None,
}
}
#[cfg(feature = "reflect_documentation")]
pub fn with_docs(self, docs: Option<&'static str>) -> Self {
Self { docs, ..self }
}
pub fn with_custom_attributes(self, custom_attributes: CustomAttributes) -> Self {
Self {
custom_attributes: Arc::new(custom_attributes),
..self
}
}
pub fn variant_names(&self) -> &[&'static str] {
&self.variant_names
}
pub fn variant(&self, name: &str) -> Option<&VariantInfo> {
self.variant_indices
.get(name)
.map(|index| &self.variants[*index])
}
pub fn variant_at(&self, index: usize) -> Option<&VariantInfo> {
self.variants.get(index)
}
pub fn index_of(&self, name: &str) -> Option<usize> {
self.variant_indices.get(name).copied()
}
pub fn variant_path(&self, name: &str) -> String {
format!("{}::{name}", self.type_path())
}
pub fn contains_variant(&self, name: &str) -> bool {
self.variant_indices.contains_key(name)
}
pub fn iter(&self) -> Iter<'_, VariantInfo> {
self.variants.iter()
}
pub fn variant_len(&self) -> usize {
self.variants.len()
}
impl_type_methods!(ty);
#[cfg(feature = "reflect_documentation")]
pub fn docs(&self) -> Option<&'static str> {
self.docs
}
impl_custom_attribute_methods!(self.custom_attributes, "enum");
impl_generic_info_methods!(generics);
}
pub struct VariantFieldIter<'a> {
container: &'a dyn Enum,
index: usize,
}
impl<'a> VariantFieldIter<'a> {
pub fn new(container: &'a dyn Enum) -> Self {
Self {
container,
index: 0,
}
}
}
impl<'a> Iterator for VariantFieldIter<'a> {
type Item = VariantField<'a>;
fn next(&mut self) -> Option<Self::Item> {
let value = match self.container.variant_type() {
VariantType::Unit => None,
VariantType::Tuple => Some(VariantField::Tuple(self.container.field_at(self.index)?)),
VariantType::Struct => {
let name = self.container.name_at(self.index)?;
Some(VariantField::Struct(name, self.container.field(name)?))
}
};
self.index += value.is_some() as usize;
value
}
fn size_hint(&self) -> (usize, Option<usize>) {
let size = self.container.field_len();
(size, Some(size))
}
}
impl<'a> ExactSizeIterator for VariantFieldIter<'a> {}
pub enum VariantField<'a> {
Struct(&'a str, &'a dyn PartialReflect),
Tuple(&'a dyn PartialReflect),
}
impl<'a> VariantField<'a> {
pub fn name(&self) -> Option<&'a str> {
if let Self::Struct(name, ..) = self {
Some(*name)
} else {
None
}
}
pub fn value(&self) -> &'a dyn PartialReflect {
match *self {
Self::Struct(_, value) | Self::Tuple(value) => value,
}
}
}
#[cfg(test)]
mod tests {
use crate::*;
#[derive(Reflect, Debug, PartialEq)]
enum MyEnum {
A,
B(usize, i32),
C { foo: f32, bar: bool },
}
#[test]
fn next_index_increment() {
let unit_enum = MyEnum::A;
let mut iter = unit_enum.iter_fields();
let size = iter.len();
for _ in 0..2 {
assert!(iter.next().is_none());
assert_eq!(size, iter.index);
}
let tuple_enum = MyEnum::B(0, 1);
let mut iter = tuple_enum.iter_fields();
let size = iter.len();
for _ in 0..2 {
let prev_index = iter.index;
assert!(iter.next().is_some());
assert_eq!(prev_index, iter.index - 1);
}
for _ in 0..2 {
assert!(iter.next().is_none());
assert_eq!(size, iter.index);
}
let struct_enum = MyEnum::C {
foo: 0.,
bar: false,
};
let mut iter = struct_enum.iter_fields();
let size = iter.len();
for _ in 0..2 {
let prev_index = iter.index;
assert!(iter.next().is_some());
assert_eq!(prev_index, iter.index - 1);
}
for _ in 0..2 {
assert!(iter.next().is_none());
assert_eq!(size, iter.index);
}
}
}