bevy_reflect: Trait casting

[MrGVSV]

Objective

Note: The majority of this PR was implemented by @kabergstrom. They kindly asked that I head up the PR due to possible unavailability to make any requested changes. But all credit should go to them!

Trait Casting

Currently the best way of converting a reflected type to a trait object is to go through a multi-step process wherein a user must:

1. Mark the trait with #[reflect_trait]
2. Import the generated ReflectSomeTrait
3. Mark the implementor structs with #[reflect(SomeTrait)]
4. Get the stored ReflectSomeTrait from the registry
5. Use the ReflectSomeTrait to cast your reflected value to &dyn SomeTrait

Example Code

#[reflect_trait]
trait SomeTrait {
  fn foo(&self);
}

#[derive(Reflect)]
#[reflect(SomeTrait)]
struct MyStruct {}

fn try_cast(value: &dyn Reflect, registry: &TypeRegistry) -> &dyn SomeTrait {
  let reflect_some_trait = registry
    .get_type_data::<ReflectSomeTrait>(value.type_id())
    .unwrap();
  
  reflect_some_trait.get(&*value).unwrap()
}

This is both complicated and annoying to code out, with lots of places to go wrong. Ideally there would be a way for casting to Just Work™.

Solution

Trait Casting

Added the following methods to TypeRegistration:

• has_trait_cast<T: ?Sized + 'static>(&self) -> bool Returns whether or not a given value can be cast to a given trait object.

  let is_implemented: bool = type_registration.has_trait_cast::<dyn SomeTrait>();
  

• trait_cast<'a, T: ?Sized + 'static>(&self, val: &'a dyn Reflect) -> Option<&'a T> Attempts to cast a given value to a given trait object.

  let obj: &dyn SomeTrait = type_registration.trait_cast::<dyn SomeTrait>(&some_struct).unwrap();
  

Together, these methods can turn the previous code:

let reflect_some_trait = registry
    .get_type_data::<ReflectSomeTrait>(value.type_id())
    .unwrap();
  
  reflect_some_trait.get(&*value).unwrap()

into the following:

let reflect_some_trait = registry
    .trait_cast::<dyn SomeTrait>(&value)
    .unwrap();

This is a small improvement, but it's not where the magic of this PR happens. That takes place during type registration.

Type Registration

The biggest difficulty with handling traits using Bevy's reflection system is that each trait has to be registered per implementor— at the implementor's definition. This can be difficult to maintain as implementors need to import the generated ReflectMyTrait structs.

To get around this issue, this PR adds the register_all macro, which allows multiple traits and types to be registered at once, and let the compiler figure out which types can register which trait.

#[derive(Reflect)]
struct Foo;
#[derive(Reflect)]
struct Bar;

trait SomeTrait {}
impl SomeTrait for Foo {}

register_all! {
  traits: [SomeTrait],
  types: [Foo, Bar]
}

// Registers:
// Foo with SomeTrait trait cast
// Bar with nothing
register_types(&mut type_registry);

The register_all registers a trait cast for each trait that a given type implements. It then generates a public function called register_types which can be used to mass register all types.

Type Data

While the focus of this PR is trait casting, we can't forget about traits which cannot be cast into dyn MyTrait objects. Those traits are often represented as type data, which can be inserted into a type's registration. Since this PR added register_all for traits, it made sense to also give the ability to register type data as well.

#[derive(Reflect)]
#[reflect(SomeTrait)]
struct Foo;
#[derive(Reflect)]
struct Bar;

#[reflect_trait]
trait SomeTrait {}
impl SomeTrait for Foo {}

register_all! {
  data: [ReflectSomeTrait],
  types: [Foo, Bar]
}

// Registers:
// Foo with ReflectSomeTrait type data
// Bar with nothing
register_types(&mut type_registry);

Note: The only required field is types, so both traits and data can be included or excluded as necessary.

Benefits

Aside from just being more compact, the macro automatically allows usage of the new trait_cast methods, making it much easier to cast types to their respective traits. It also makes using managing reflected traits a lot simpler (users don't need to worry about importing/exporting the proper ReflectMyTrait structs).

Additionally, it makes it much easier to reflect third-party traits. Instead of manually creating the struct (or using a macro to do it for you), you can simply add it to the list of traits in the register_all! macro (same goes for third-party structs as well).

Considerations

One downside of this is that it ends up generating a fair amount of code (~40 lines per type per trait) as each type needs to check if it implements a given trait. This means that the total number of generated "blocks" is equal to the number of types times the number of traits.

However, it might be worth it for widely used traits or when used more granularly.

Sample Output

pub fn register_types(registry: &mut bevy_reflect::TypeRegistry) {
  {
    let type_registration = match registry.get_mut(::std::any::TypeId::of::<Foo>()) {
      Some(registration) => registration,
      None => {
        registry.register::<Foo>();
        registry.get_mut(::std::any::TypeId::of::<Foo>()).unwrap()
      }
    };
    if let Some(cast_fn) = {
      {
        trait NotTrait {
          const CAST_FN: Option<for<'a> fn(&'a Foo) -> &'a dyn SomeTrait> = None;
        }
        impl<T> NotTrait for T {}
        struct IsImplemented<T>(core::marker::PhantomData<T>);
        impl<T: SomeTrait + 'static> IsImplemented<T> {
          #[allow(dead_code)]
          const CAST_FN: Option<for<'a> fn(&'a T) -> &'a dyn SomeTrait> =
            Some(|a| a);
        }
        if IsImplemented::<Foo>::CAST_FN.is_some() {
          let f: fn(&dyn crate::Reflect) -> crate::ErasedNonNull =
            |val: &dyn crate::Reflect| {
              let cast_fn = IsImplemented::<Foo>::CAST_FN.unwrap();
              let static_val: &Foo = val.downcast_ref::<Foo>().unwrap();
              let trait_val: &dyn SomeTrait = (cast_fn)(static_val);
              crate::ErasedNonNull::new(
                trait_val,
                core::any::TypeId::of::<dyn SomeTrait>(),
              )
            };
          Some(f)
        } else {
          None
        }
      }
    } {
      {
        type_registration.register_trait_cast::<dyn SomeTrait>(cast_fn);
      }
    }
  };
  {
    let type_registration = match registry.get_mut(::std::any::TypeId::of::<Bar>()) {
      Some(registration) => registration,
      None => {
        registry.register::<Bar>();
        registry.get_mut(::std::any::TypeId::of::<Bar>()).unwrap()
      }
    };
    if let Some(cast_fn) = {
      {
        trait NotTrait {
          const CAST_FN: Option<for<'a> fn(&'a Bar) -> &'a dyn SomeTrait> = None;
        }
        impl<T> NotTrait for T {}
        struct IsImplemented<T>(core::marker::PhantomData<T>);
        impl<T: SomeTrait + 'static> IsImplemented<T> {
          #[allow(dead_code)]
          const CAST_FN: Option<for<'a> fn(&'a T) -> &'a dyn SomeTrait> =
            Some(|a| a);
        }
        if IsImplemented::<Bar>::CAST_FN.is_some() {
          let f: fn(&dyn crate::Reflect) -> crate::ErasedNonNull =
            |val: &dyn crate::Reflect| {
              let cast_fn = IsImplemented::<Bar>::CAST_FN.unwrap();
              let static_val: &Bar = val.downcast_ref::<Bar>().unwrap();
              let trait_val: &dyn SomeTrait = (cast_fn)(static_val);
              crate::ErasedNonNull::new(
                trait_val,
                core::any::TypeId::of::<dyn SomeTrait>(),
              )
            };
          Some(f)
        } else {
          None
        }
      }
    } {
      {
        type_registration.register_trait_cast::<dyn SomeTrait>(cast_fn);
      }
    }
  }
}

Example

As an example of how this works, bevy_reflect::impls now exports a registrations module:

mod registrations {
    use crate as bevy_reflect;
    use crate::erased_serde::Serialize;
    use crate::register_all;

    register_all! {
        traits: [Serialize],
        types: [
            bool,
            char,
            u8,
            u16,
            u32,
            u64,
            u128,
            usize,
            i8,
            i16,
            i32,
            i64,
            i128,
            isize,
            f32,
            f64,
            String,
            Option<String>,
        ]
    }
}

Because we register all these types with a Serialize trait cast, we can use it in our ReflectSerializer as easily as:

// `unwrap()` used for conciseness here
let registration = self.registry
  .get(self.value.type_id())
  .unwrap();
let serializable = registration
  .trait_cast::<dyn erased_serde::Serialize>(self.value)
  .unwrap();

[PROMETHIA-27]

I'm also not a fan of register_all, it's really not the direction I like to see any library move in. Why can't this be a part of the #[reflect(Trait)] syntax?

[kabergstrom]

I'm also not a fan of register_all, it's really not the direction I like to see any library move in. Why can't this be a part of the #[reflect(Trait)] syntax?

Many traits are not implemented for a struct in the same crate as the struct is declared.

[kabergstrom]

Finally giving this the attention it deserves. My current thoughts:

1. The trait casting improvements are welcome and great ergonomics.

2. Something like the register_all! macro is essential to reduce boilerplate, but it's syntax feels alien to me and it's hard to tell what's going on. Generating a named function that you must call definitely feels wrong. Can we do something method-like on type registeries? Can we pass in a type registry to the macro?

3. Why have all the `Serialize` reflect_value calls been removed but not the others?

Regarding 2, I suppose we could require the user to make their own function, and then they invoke the macro inside of it to expand the code.

[MrGVSV]

2. Something like the register_all! macro is essential to reduce boilerplate, but it's syntax feels alien to me and it's hard to tell what's going on. Generating a named function that you must call definitely feels wrong. Can we do something method-like on type registeries? Can we pass in a type registry to the macro?

Yeah that's totally fair. What if it looked more like:

let mut app = App::new();
register_all! {
  app,
  traits: [SomeTrait],
  types: [Foo, Bar]
};
// ...

Where we instead pass the app in directly? Or maybe it should be the registry 🤔

3. Why have all the Serialize reflect_value calls been removed but not the others?

The reason is that, with the current serialization system, only ReflectRef::Value types are actually use the impl'd Serialize. Otherwise, they always use the reflected form. This is addressed in #5723.

---

I'm also not a fan of register_all, it's really not the direction I like to see any library move in. Why can't this be a part of the #[reflect(Trait)] syntax?

I agree. I think register_all! is great for mass registration, keeping everything in one place, and registering for foreign types. However, I think whatever we do, the #[reflect(SomeTrait)] syntax needs to be updated to support the trait casting registration (maybe as #[reflect(dyn SomeTrait)]?).

[alice-i-cecile]

Where we instead pass the app in directly? Or maybe it should be the registry 🤔

I like this much better! I think it should be the registry though.

(maybe as #[reflect(dyn SomeTrait)]?).

This syntax makes sense to me. We'll need good docs to support it though!

[MrGVSV]

Closing this out as stale. I also don't know if it's the direction we want to go.