rbx_binary: Reuse Descriptors

rbx_binary/src/core.rs

@@ -3,7 +3,6 @@ use std::{
     mem,
 };
 
-use rbx_dom_weak::Ustr;
 use rbx_reflection::{
     ClassDescriptor, PropertyDescriptor, PropertyKind, PropertySerialization, ReflectionDatabase,
 };
@@ -342,99 +341,95 @@ pub struct PropertyDescriptors<'db> {
     pub serialized: Option<&'db PropertyDescriptor<'db>>,
 }
 
-/// Find both the canonical and serialized property descriptors for a given
-/// class and property name pair. These might be the same descriptor!
-pub fn find_property_descriptors<'db>(
-    database: &'db ReflectionDatabase<'db>,
-    class_name: Ustr,
-    property_name: Ustr,
-) -> Option<PropertyDescriptors<'db>> {
-    let mut class_descriptor = database.classes.get(class_name.as_str())?;
+impl<'db> PropertyDescriptors<'db> {
+    /// Get both the canonical and serialized property descriptors for a given
+    /// class and property descriptor. The canonical and serialized descriptors
+    /// might be the same descriptor!
+    pub fn new(
+        class_descriptor: &'db ClassDescriptor<'db>,
+        property_descriptor: &'db PropertyDescriptor<'db>,
+    ) -> Option<PropertyDescriptors<'db>> {
+        match &property_descriptor.kind {
+            // This property descriptor is the canonical form of this
+            // logical property.
+            PropertyKind::Canonical { serialization } => {
+                let serialized = find_serialized_from_canonical(
+                    class_descriptor,
+                    property_descriptor,
+                    serialization,
+                );
+
+                Some(PropertyDescriptors {
+                    canonical: property_descriptor,
+                    serialized,
+                })
+            }
 
-    // We need to find the canonical property descriptor associated with
-    // the property we're working with.
-    //
-    // At each step of the loop, we're checking a new class descriptor to see if
-    // it has an entry for the property name we're looking for. If that class
-    // doesn't have the property, we'll check its superclass until we reach the
-    // root.
-    loop {
-        // If this class descriptor knows about this property name, we're pretty
-        // much done!
-        if let Some(property_descriptor) = class_descriptor.properties.get(property_name.as_str()) {
-            match &property_descriptor.kind {
-                // This property descriptor is the canonical form of this
-                // logical property. That means we've found one of the two
-                // descriptors we're looking for!
-                PropertyKind::Canonical { serialization } => {
-                    let serialized = find_serialized_from_canonical(
-                        class_descriptor,
-                        property_descriptor,
-                        serialization,
-                    );
+            // This descriptor is an alias for another property. While this
+            // descriptor might be one of the two descriptors we need to
+            // return, it's possible that both the canonical and serialized
+            // forms are different.
+            PropertyKind::Alias { alias_for } => {
+                let canonical = class_descriptor.properties.get(alias_for.as_ref()).unwrap();
 
-                    return Some(PropertyDescriptors {
-                        canonical: property_descriptor,
+                if let PropertyKind::Canonical { serialization } = &canonical.kind {
+                    let serialized =
+                        find_serialized_from_canonical(class_descriptor, canonical, serialization);
+
+                    Some(PropertyDescriptors {
+                        canonical,
                         serialized,
-                    });
-                }
+                    })
+                } else {
+                    // If one property in the database calls itself an alias
+                    // of another property, that property must be canonical.
+                    log::error!(
+                        "Property {}.{} is marked as an alias for {}.{}, but the latter is not canonical.",
+                        class_descriptor.name,
+                        property_descriptor.name,
+                        class_descriptor.name,
+                        alias_for
+                    );
 
-                // This descriptor is an alias for another property. While this
-                // descriptor might be one of the two descriptors we need to
-                // return, it's possible that both the canonical and serialized
-                // forms are different.
-                PropertyKind::Alias { alias_for } => {
-                    let canonical = class_descriptor.properties.get(alias_for.as_ref()).unwrap();
-
-                    if let PropertyKind::Canonical { serialization } = &canonical.kind {
-                        let serialized = find_serialized_from_canonical(
-                            class_descriptor,
-                            canonical,
-                            serialization,
-                        );
-
-                        return Some(PropertyDescriptors {
-                            canonical,
-                            serialized,
-                        });
-                    } else {
-                        // If one property in the database calls itself an alias
-                        // of another property, that property must be canonical.
-                        log::error!(
-                            "Property {}.{} is marked as an alias for {}.{}, but the latter is not canonical.",
-                            class_descriptor.name,
-                            property_descriptor.name,
-                            class_descriptor.name,
-                            alias_for
-                        );
-
-                        return None;
-                    }
+                    None
                 }
-
-                // This descriptor is of an unknown kind and we don't know how
-                // to deal with it -- maybe rbx_binary is out of date?
-                _ => return None,
             }
-        }
-
-        if let Some(superclass_name) = &class_descriptor.superclass {
-            // If a property descriptor isn't found in our class, check our
-            // superclass.
-
-            class_descriptor = database
-                .classes
-                .get(superclass_name)
-                .expect("Superclass in reflection database didn't exist");
-        } else {
-            // This property isn't known by any class in the reflection
-            // database.
 
-            return None;
+            // This descriptor is of an unknown kind and we don't know how
+            // to deal with it -- maybe rbx_binary is out of date?
+            _ => None,
         }
     }
 }
 
+/// Find the superclass which contains the specified property,
+/// extract the canonical and serialized property descriptors,
+/// and return both.
+pub fn find_property_descriptors<'db>(
+    database: &'db ReflectionDatabase<'db>,
+    class_descriptor: Option<&'db ClassDescriptor<'db>>,
+    property_name: &str,
+) -> Option<(&'db ClassDescriptor<'db>, PropertyDescriptors<'db>)> {
+    // Checking the class descriptor is ugly without an optional
+    // return value, and all the call sites need this precise logic.
+    let class_descriptor = class_descriptor?;
+
+    // We need to find the canonical property descriptor associated with
+    // the property we're working with. Walk superclasses and
+    // find a class descriptor which knows about this property name.
+    let (class, prop) = database
+        .superclasses_iter(class_descriptor)
+        .find_map(|class| {
+            let prop = class.properties.get(property_name)?;
+            Some((class, prop))
+        })?;
+
+    // Extract the canonical and serialized property descriptors
+    // from the class and property descriptors
+    let descriptors = PropertyDescriptors::new(class, prop)?;
+    Some((class, descriptors))
+}
+
 /// Given the canonical property descriptor for a logical property along with
 /// its serialization, returns the serialized form of the logical property if
 /// this property is serializable.

rbx_binary/src/deserializer/state.rs

@@ -12,7 +12,9 @@ use rbx_dom_weak::{
     },
     InstanceBuilder, Ustr, WeakDom,
 };
-use rbx_reflection::{DataType, PropertyKind, PropertySerialization, ReflectionDatabase};
+use rbx_reflection::{
+    ClassDescriptor, DataType, PropertyKind, PropertySerialization, ReflectionDatabase,
+};
 
 use crate::{
     chunk::Chunk,
@@ -42,7 +44,7 @@ pub(super) struct DeserializerState<'db, R> {
     shared_strings: Vec<SharedString>,
 
     /// All of the instance types described by the file so far.
-    type_infos: HashMap<u32, TypeInfo>,
+    type_infos: HashMap<u32, TypeInfo<'db>>,
 
     /// All of the instances known by the deserializer.
     instances_by_ref: HashMap<i32, Instance>,
@@ -59,7 +61,7 @@ pub(super) struct DeserializerState<'db, R> {
 
 /// Represents a unique instance class. Binary models define all their instance
 /// types up front and give them a short u32 identifier.
-struct TypeInfo {
+struct TypeInfo<'db> {
     /// The ID given to this type by the current file we're deserializing. This
     /// ID can be different for different files.
     type_id: u32,
@@ -69,6 +71,10 @@ struct TypeInfo {
 
     /// A list of the instances described by this file that are this type.
     referents: Vec<i32>,
+
+    /// A reference to the type's class descriptor from rbx_reflection, if this
+    /// is a known class.
+    class_descriptor: Option<&'db ClassDescriptor<'db>>,
 }
 
 /// Contains all the information we need to gather in order to construct an
@@ -97,11 +103,11 @@ struct CanonicalProperty<'db> {
 fn find_canonical_property<'de>(
     database: &'de ReflectionDatabase,
     binary_type: Type,
-    class_name: Ustr,
-    prop_name: Ustr,
+    class_descriptor: Option<&'de ClassDescriptor<'de>>,
+    prop_name: &str,
 ) -> Option<CanonicalProperty<'de>> {
-    match find_property_descriptors(database, class_name, prop_name) {
-        Some(descriptors) => {
+    match find_property_descriptors(database, class_descriptor, prop_name) {
+        Some((_, descriptors)) => {
             // If this descriptor is known but wasn't supposed to be
             // serialized, we should skip it.
             //
@@ -169,7 +175,7 @@ fn find_canonical_property<'de>(
             log::trace!("Unknown prop, using type {:?}", canonical_type);
 
             Some(CanonicalProperty {
-                name: prop_name,
+                name: prop_name.into(),
                 ty: canonical_type,
                 migration: None,
             })
@@ -291,13 +297,12 @@ impl<'db, R: Read> DeserializerState<'db, R> {
         let mut referents = vec![0; number_instances as usize];
         chunk.read_referent_array(&mut referents)?;
 
-        let prop_capacity = self
-            .deserializer
-            .database
-            .classes
-            .get(type_name.as_str())
-            .map(|class| class.default_properties.len())
-            .unwrap_or(0);
+        let (class_descriptor, prop_capacity) =
+            if let Some(class) = self.deserializer.database.classes.get(type_name.as_str()) {
+                (Some(class), class.default_properties.len())
+            } else {
+                (None, 0)
+            };
 
         // TODO: Check object_format and check for service markers if it's 1?
 
@@ -320,6 +325,7 @@ impl<'db, R: Read> DeserializerState<'db, R> {
                 type_id,
                 type_name: type_name.into(),
                 referents,
+                class_descriptor,
             },
         );
 
@@ -405,8 +411,8 @@ This may cause unexpected or broken behavior in your final results if you rely o
         let property = if let Some(property) = find_canonical_property(
             self.deserializer.database,
             binary_type,
-            type_info.type_name,
-            prop_name.as_str().into(),
+            type_info.class_descriptor,
+            &prop_name,
         ) {
             property
         } else {

rbx_binary/src/serializer/state.rs

@@ -25,7 +25,8 @@ use rbx_reflection::{
 use crate::{
     chunk::ChunkBuilder,
     core::{
-        find_property_descriptors, RbxWriteExt, FILE_MAGIC_HEADER, FILE_SIGNATURE, FILE_VERSION,
+        find_property_descriptors, PropertyDescriptors, RbxWriteExt, FILE_MAGIC_HEADER,
+        FILE_SIGNATURE, FILE_VERSION,
     },
     types::Type,
     Serializer,
@@ -345,8 +346,8 @@ impl<'dom, 'db, W: Write> SerializerState<'dom, 'db, W> {
             let mut migration = None;
 
             let database = self.serializer.database;
-            match find_property_descriptors(database, instance.class, *prop_name) {
-                Some(descriptors) => {
+            match find_property_descriptors(database, type_info.class_descriptor, prop_name) {
+                Some((superclass_descriptor, descriptors)) => {
                     // For any properties that do not serialize, we can skip
                     // adding them to the set of type_infos.
                     let serialized = match descriptors.serialized {
@@ -360,11 +361,16 @@ impl<'dom, 'db, W: Write> SerializerState<'dom, 'db, W> {
                                 // information of the new property instead of the old
                                 // property, because migrated properties should not
                                 // serialize
-                                let new_descriptors = find_property_descriptors(
-                                    database,
-                                    instance.class,
-                                    prop_migration.new_property_name.as_str().into(),
-                                );
+                                //
+                                // Assume that the migration will always be
+                                // directed to a property on the same class.
+                                // This avoids re-walking the superclasses.
+                                let new_descriptors = superclass_descriptor
+                                    .properties
+                                    .get(prop_migration.new_property_name.as_str())
+                                    .and_then(|prop| {
+                                        PropertyDescriptors::new(superclass_descriptor, prop)
+                                    });
 
                                 migration = Some(prop_migration);