bevy_reflect: enforce that type data is only inserted for the correct type

crates/bevy_reflect/bevy_reflect_derive/src/lib.rs

@@ -719,7 +719,7 @@ fn impl_get_type_registration(
         impl #impl_generics #bevy_reflect_path::GetTypeRegistration for #type_name #ty_generics #where_clause {
             fn get_type_registration() -> #bevy_reflect_path::TypeRegistration {
                 let mut registration = #bevy_reflect_path::TypeRegistration::of::<#type_name #ty_generics>();
-                #(registration.insert::<#registration_data>(#bevy_reflect_path::FromType::<#type_name #ty_generics>::from_type());)*
+                #(registration.insert::<Self, #registration_data>();)*
                 registration
             }
         }

crates/bevy_reflect/src/impls/std.rs

@@ -1,8 +1,7 @@
 use crate as bevy_reflect;
 use crate::{
-    map_partial_eq, serde::Serializable, DynamicMap, FromReflect, FromType, GetTypeRegistration,
-    List, ListIter, Map, MapIter, Reflect, ReflectDeserialize, ReflectMut, ReflectRef,
-    TypeRegistration,
+    map_partial_eq, serde::Serializable, DynamicMap, FromReflect, GetTypeRegistration, List,
+    ListIter, Map, MapIter, Reflect, ReflectDeserialize, ReflectMut, ReflectRef, TypeRegistration,
 };
 
 use bevy_reflect_derive::{impl_from_reflect_value, impl_reflect_value};
@@ -147,7 +146,7 @@ unsafe impl<T: FromReflect> Reflect for Vec<T> {
 impl<T: FromReflect + for<'de> Deserialize<'de>> GetTypeRegistration for Vec<T> {
     fn get_type_registration() -> TypeRegistration {
         let mut registration = TypeRegistration::of::<Vec<T>>();
-        registration.insert::<ReflectDeserialize>(FromType::<Vec<T>>::from_type());
+        registration.insert::<Self, ReflectDeserialize>();
         registration
     }
 }
@@ -269,7 +268,7 @@ where
 {
     fn get_type_registration() -> TypeRegistration {
         let mut registration = TypeRegistration::of::<Self>();
-        registration.insert::<ReflectDeserialize>(FromType::<Self>::from_type());
+        registration.insert::<Self, ReflectDeserialize>();
         registration
     }
 }
@@ -354,7 +353,7 @@ unsafe impl Reflect for Cow<'static, str> {
 impl GetTypeRegistration for Cow<'static, str> {
     fn get_type_registration() -> TypeRegistration {
         let mut registration = TypeRegistration::of::<Cow<'static, str>>();
-        registration.insert::<ReflectDeserialize>(FromType::<Cow<'static, str>>::from_type());
+        registration.insert::<Self, ReflectDeserialize>();
         registration
     }
 }

crates/bevy_reflect/src/type_registry.rs

@@ -182,27 +182,30 @@ impl TypeRegistration {
     /// data.
     ///
     /// Returns `None` if no such value exists.
+    ///
+    /// The value is guaranteed to be a value of `T` generated using [`FromType`] of the type that this
+    /// registration is used for.
     pub fn data<T: TypeData>(&self) -> Option<&T> {
         self.data
             .get(&TypeId::of::<T>())
             .and_then(|value| value.downcast_ref())
     }
 
-    /// Returns a mutable reference to the value of type `T` in this
-    /// registration's type data.
+    /// Inserts an instance of `D` into this registration's type data.
+    /// The value is created using the [`FromType<T>`] impl for `D`
+    /// and `T` must match the type that this [`TypeRegistration`] was created for.
     ///
-    /// Returns `None` if no such value exists.
-    pub fn data_mut<T: TypeData>(&mut self) -> Option<&mut T> {
-        self.data
-            .get_mut(&TypeId::of::<T>())
-            .and_then(|value| value.downcast_mut())
-    }
-
-    /// Inserts an instance of `T` into this registration's type data.
-    ///
-    /// If another instance of `T` was previously inserted, it is replaced.
-    pub fn insert<T: TypeData>(&mut self, data: T) {
-        self.data.insert(TypeId::of::<T>(), Box::new(data));
+    /// If another instance of `D` was previously inserted, it is replaced.
+    pub fn insert<T: 'static, D: TypeData + FromType<T>>(&mut self) {
+        if self.type_id != std::any::TypeId::of::<T>() {
+            panic!(
+            "called `TypeRegistration::insert` on a registration for `{}` with data for type `{}`",
+            self.short_name,
+            std::any::type_name::<T>()
+        );
+        }
+        let data = <D as FromType<T>>::from_type();
+        self.data.insert(TypeId::of::<D>(), Box::new(data));
     }
 
     /// Creates type registration information for `T`.
@@ -308,8 +311,8 @@ where
 
 /// Trait used to generate [`TypeData`] for trait reflection.
 ///
-/// This is used by the `#[derive(Reflect)]` macro to generate an implementation
-/// of [`TypeData`] to pass to [`TypeRegistration::insert`].
+/// This is used by [`TypeRegistration::insert`] to generate an implementation
+/// of [`TypeData`] for the type `T`.
 pub trait FromType<T> {
     fn from_type() -> Self;
 }
@@ -352,7 +355,7 @@ impl<T: for<'a> Deserialize<'a> + Reflect> FromType<T> for ReflectDeserialize {
 
 #[cfg(test)]
 mod test {
-    use crate::TypeRegistration;
+    use crate::{FromType, TypeRegistration};
     use bevy_utils::HashMap;
 
     #[test]
@@ -423,4 +426,27 @@ mod test {
             "Option<HashMap<Option<String>, (String, Option<String>)>>"
         );
     }
+
+    #[derive(Clone)]
+    struct SomeTypeData;
+    impl<T> FromType<T> for SomeTypeData {
+        fn from_type() -> Self {
+            SomeTypeData
+        }
+    }
+
+    #[test]
+    fn test_type_registration_insert() {
+        let mut registration = TypeRegistration::of::<f64>();
+        registration.insert::<f64, SomeTypeData>();
+
+        assert!(registration.data::<SomeTypeData>().is_some());
+    }
+
+    #[test]
+    #[should_panic = "for `f64` with data for type `f32`"]
+    fn test_type_registration_wrong_type() {
+        let mut registration = TypeRegistration::of::<f64>();
+        registration.insert::<f32, SomeTypeData>();
+    }
 }