Rust eager serialization + partial updates playground

crates/build/re_types_builder/src/codegen/rust/api.rs

@@ -25,8 +25,8 @@ use crate::{
     ArrowRegistry, CodeGenerator, ElementType, Object, ObjectField, ObjectKind, Objects, Reporter,
     Type, ATTR_DEFAULT, ATTR_RERUN_COMPONENT_OPTIONAL, ATTR_RERUN_COMPONENT_RECOMMENDED,
     ATTR_RERUN_COMPONENT_REQUIRED, ATTR_RERUN_LOG_MISSING_AS_EMPTY, ATTR_RERUN_VIEW_IDENTIFIER,
-    ATTR_RUST_CUSTOM_CLAUSE, ATTR_RUST_DERIVE, ATTR_RUST_DERIVE_ONLY, ATTR_RUST_NEW_PUB_CRATE,
-    ATTR_RUST_REPR,
+    ATTR_RUST_ARCHETYPE_EAGER, ATTR_RUST_CUSTOM_CLAUSE, ATTR_RUST_DERIVE, ATTR_RUST_DERIVE_ONLY,
+    ATTR_RUST_NEW_PUB_CRATE, ATTR_RUST_REPR,
 };
 
 use super::{
@@ -233,7 +233,15 @@ fn generate_mod_file(
     {
         let module_name = obj.snake_case_name();
         let type_name = &obj.name;
-        code.push_str(&format!("pub use self::{module_name}::{type_name};\n"));
+        let native_type_name = format!("Native{type_name}");
+
+        if obj.requires_native_rust_archetype() {
+            code.push_str(&format!(
+                "pub use self::{module_name}::{{{type_name}, {native_type_name}}};\n"
+            ));
+        } else {
+            code.push_str(&format!("pub use self::{module_name}::{type_name};\n"));
+        }
     }
     // And then deprecated.
     if objects.iter().any(|obj| obj.deprecation_notice().is_some()) {
@@ -245,10 +253,19 @@ fn generate_mod_file(
     {
         let module_name = obj.snake_case_name();
         let type_name = &obj.name;
+        let native_type_name = format!("Native{type_name}");
+
         if obj.deprecation_notice().is_some() {
             code.push_str("#[allow(deprecated)]\n");
         }
-        code.push_str(&format!("pub use self::{module_name}::{type_name};\n"));
+
+        if obj.requires_native_rust_archetype() {
+            code.push_str(&format!(
+                "pub use self::{module_name}::{{{type_name}, {native_type_name}}};\n"
+            ));
+        } else {
+            code.push_str(&format!("pub use self::{module_name}::{type_name};\n"));
+        }
     }
 
     files_to_write.insert(path, code);
@@ -264,6 +281,15 @@ fn quote_struct(
 ) -> TokenStream {
     assert!(obj.is_struct());
 
+    // Certain eager archetypes might require the generation of an associated native archetype, as
+    // the internal viewer code heavily relies on it.
+    let obj_native = obj.requires_native_rust_archetype().then(|| {
+        let mut obj_native = obj.clone();
+        obj_native.name = format!("Native{}", obj_native.name);
+        obj_native.attrs.remove(ATTR_RUST_ARCHETYPE_EAGER);
+        obj_native
+    });
+
     let Object { name, fields, .. } = obj;
 
     let name = format_ident!("{name}");
@@ -309,10 +335,74 @@ fn quote_struct(
     } else {
         quote! { pub struct #name { #(#quoted_fields,)* }}
     };
+    let quoted_struct_native = obj_native.as_ref().map(|obj_native| {
+        let native_name = format_ident!("{}", obj_native.name);
+
+        let quoted_fields = obj_native
+            .fields
+            .iter()
+            .map(|obj_field| ObjectFieldTokenizer(reporter, obj_native, obj_field).quoted(objects));
+        let quoted_struct = if is_tuple_struct {
+            quote! { pub struct #native_name(#(#quoted_fields,)*); }
+        } else {
+            quote! { pub struct #native_name { #(#quoted_fields,)* }}
+        };
+
+        let eager_fields_to_native_fields = obj.fields.iter().map(|field| {
+            let field_name = format_ident!("{}", field.name);
+            quote!(value.#field_name.clone().map(|batch| (batch.descriptor, batch.array)))
+        });
+        let eager_to_native = quote! {
+            impl TryFrom<&#name> for #native_name {
+                type Error = crate::DeserializationError;
+
+                #[rustfmt::skip] // so it doesn't take 1000 lines for no reason ><
+                fn try_from(value: &#name) -> Result<Self, Self::Error> {
+                    use ::re_types_core::Archetype as _;
+                    Self::from_arrow_components(
+                        [ #(#eager_fields_to_native_fields),* ]
+                        .into_iter()
+                        .flatten(),
+                    )
+                }
+            }
+        };
+
+        let native_fields_to_eager_fields = obj_native.fields.iter().map(|field| {
+            let field_name = format_ident!("{}", field.name);
+            if field.is_nullable {
+                quote!(#field_name: value.#field_name.as_ref().and_then(|v| v.serialized()))
+            } else {
+                quote!(#field_name: value.#field_name.serialized())
+            }
+        });
+        let native_to_eager = quote! {
+            impl From<&#native_name> for #name {
+                #[rustfmt::skip] // so it doesn't take 1000 lines for no reason ><
+                #[inline]
+                fn from(value: &#native_name) -> Self {
+                    Self {
+                        #(#native_fields_to_eager_fields),*
+                    }
+                }
+            }
+        };
+
+        quote! {
+            #[doc(hidden)]
+            #quoted_derive_clone_debug
+            #quoted_struct
+
+            #eager_to_native
+
+            #native_to_eager
+        }
+    });
 
     let quoted_from_impl = quote_from_impl_from_obj(obj);
 
-    let quoted_trait_impls = quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj);
+    let quoted_trait_impls =
+        quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj, obj_native.as_ref());
 
     let quoted_builder = quote_builder_from_obj(reporter, objects, obj);
 
@@ -368,6 +458,8 @@ fn quote_struct(
         #quoted_deprecation_notice
         #quoted_struct
 
+        #quoted_struct_native
+
         #quoted_trait_impls
 
         #quoted_from_impl
@@ -426,7 +518,8 @@ fn quote_union(
         }
     });
 
-    let quoted_trait_impls = quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj);
+    let quoted_trait_impls =
+        quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj, None);
 
     let quoted_heap_size_bytes = {
         let quoted_matches = fields.iter().map(|obj_field| {
@@ -574,7 +667,8 @@ fn quote_enum(
         }
     });
 
-    let quoted_trait_impls = quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj);
+    let quoted_trait_impls =
+        quote_trait_impls_from_obj(reporter, arrow_registry, objects, obj, None);
 
     let all = fields.iter().map(|field| {
         let name = format_ident!("{}", field.name);
@@ -843,13 +937,14 @@ fn quote_trait_impls_from_obj(
     arrow_registry: &ArrowRegistry,
     objects: &Objects,
     obj: &Object,
+    obj_native: Option<&Object>,
 ) -> TokenStream {
     match obj.kind {
         ObjectKind::Datatype | ObjectKind::Component => {
             quote_trait_impls_for_datatype_or_component(objects, arrow_registry, obj)
         }
 
-        ObjectKind::Archetype => quote_trait_impls_for_archetype(obj),
+        ObjectKind::Archetype => quote_trait_impls_for_archetype(obj, obj_native),
 
         ObjectKind::View => quote_trait_impls_for_view(reporter, obj),
     }
@@ -1049,7 +1144,7 @@ fn quote_trait_impls_for_datatype_or_component(
     }
 }
 
-fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
+fn quote_trait_impls_for_archetype(obj: &Object, obj_native: Option<&Object>) -> TokenStream {
     #![allow(clippy::collapsible_else_if)]
 
     let Object {
@@ -1260,7 +1355,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
     };
 
     // TODO(#7245): This goes away once all archetypes have been made eager.
-    let all_native_deserializers = {
+    let all_native_deserializers = |origin: TokenStream| {
         obj.fields.iter().map(|obj_field| {
             let obj_field_fqname = obj_field.fqname.as_str();
             let field_name = format_ident!("{}", obj_field.name);
@@ -1289,7 +1384,6 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
                 }
             };
 
-
             // NOTE: An archetype cannot have overlapped component types by definition, so use the
             // component's fqname to do the mapping.
             let quoted_deser = if is_nullable && !is_plural {
@@ -1303,7 +1397,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
                     };
 
                 quote! {
-                    if let Some(array) = arrays_by_descr.get(&Self::#descr_fn_name()) {
+                    if let Some(array) = arrays_by_descr.get(&#origin::#descr_fn_name()) {
                         <#component>::from_arrow_opt(&**array)
                             .with_context(#obj_field_fqname)?
                             #quoted_collection
@@ -1313,7 +1407,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
                 }
             } else if is_nullable {
                 quote! {
-                    if let Some(array) = arrays_by_descr.get(&Self::#descr_fn_name()) {
+                    if let Some(array) = arrays_by_descr.get(&#origin::#descr_fn_name()) {
                         Some({
                             <#component>::from_arrow_opt(&**array)
                                 .with_context(#obj_field_fqname)?
@@ -1326,7 +1420,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
             } else {
                 quote! {{
                     let array = arrays_by_descr
-                        .get(&Self::#descr_fn_name())
+                        .get(&#origin::#descr_fn_name())
                         .ok_or_else(DeserializationError::missing_data)
                         .with_context(#obj_field_fqname)?;
 
@@ -1335,7 +1429,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
             };
 
             quote!(let #field_name = #quoted_deser;)
-        })
+        }).collect_vec()
     };
 
     let all_eager_deserializers = {
@@ -1356,9 +1450,33 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
     let all_deserializers = if obj.is_eager_rust_archetype() {
         quote!(#(#all_eager_deserializers;)*)
     } else {
+        let all_native_deserializers = all_native_deserializers(quote!(Self));
         quote!(#(#all_native_deserializers;)*)
     };
 
+    let from_arrow_components_native = obj_native.map(|obj_native| {
+        let native_name = format_ident!("{}", obj_native.name);
+
+        let all_native_deserializers = all_native_deserializers(quote!(#name));
+        quote! {
+            impl #native_name {
+                fn from_arrow_components(
+                    arrow_data: impl IntoIterator<Item = (ComponentDescriptor, arrow::array::ArrayRef)>,
+                ) -> DeserializationResult<Self> {
+                    re_tracing::profile_function!();
+                    use ::re_types_core::{Loggable as _, ResultExt as _};
+
+                    let arrays_by_descr: ::nohash_hasher::IntMap<_, _> = arrow_data.into_iter().collect();
+                    #(#all_native_deserializers;)*
+
+                    Ok(Self {
+                        #(#quoted_field_names,)*
+                    })
+                }
+            }
+        }
+    });
+
     quote! {
         impl #name {
             #(#all_descriptor_methods)*
@@ -1436,7 +1554,6 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
                 use ::re_types_core::{Loggable as _, ResultExt as _};
 
                 let arrays_by_descr: ::nohash_hasher::IntMap<_, _> = arrow_data.into_iter().collect();
-
                 #all_deserializers
 
                 Ok(Self {
@@ -1445,6 +1562,8 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
             }
         }
 
+        #from_arrow_components_native
+
         impl ::re_types_core::AsComponents for #name {
             #as_components_impl
         }
@@ -1496,7 +1615,7 @@ fn quote_from_impl_from_obj(obj: &Object) -> TokenStream {
     let self_field_access = if obj_is_tuple_struct {
         quote!(self.0)
     } else {
-        quote!(self.#quoted_obj_field_name )
+        quote!(self.#quoted_obj_field_name)
     };
     let deref_impl = quote! {
         impl std::ops::Deref for #quoted_obj_name {

crates/build/re_types_builder/src/codegen/rust/util.rs

@@ -133,6 +133,8 @@ pub fn string_from_quoted(
 
         let line_is_attr = trimmed.starts_with("#[allow(")
             || trimmed.starts_with("#[inline]")
+            || trimmed.starts_with("#[doc(hidden)]")
+            || trimmed.starts_with("#[rustfmt::skip]")
             || trimmed.starts_with("#[derive");
 
         if line_is_attr && (!prev_line_was_attr && !prev_line_was_docstring) {

crates/build/re_types_builder/src/lib.rs

@@ -198,6 +198,7 @@ pub const ATTR_CPP_NO_FIELD_CTORS: &str = "attr.cpp.no_field_ctors";
 pub const ATTR_CPP_RENAME_FIELD: &str = "attr.cpp.rename_field";
 
 pub const ATTR_RUST_ARCHETYPE_EAGER: &str = "attr.rust.archetype_eager";
+pub const ATTR_RUST_ARCHETYPE_NATIVE: &str = "attr.rust.archetype_native";
 pub const ATTR_RUST_CUSTOM_CLAUSE: &str = "attr.rust.custom_clause";
 pub const ATTR_RUST_DERIVE: &str = "attr.rust.derive";
 pub const ATTR_RUST_DERIVE_ONLY: &str = "attr.rust.derive_only";

crates/build/re_types_builder/src/objects.rs

@@ -13,6 +13,7 @@ use crate::{
     root_as_schema, Docs, FbsBaseType, FbsEnum, FbsEnumVal, FbsField, FbsKeyValue, FbsObject,
     FbsSchema, FbsType, Reporter, ATTR_RERUN_COMPONENT_OPTIONAL, ATTR_RERUN_COMPONENT_RECOMMENDED,
     ATTR_RERUN_COMPONENT_REQUIRED, ATTR_RERUN_OVERRIDE_TYPE, ATTR_RUST_ARCHETYPE_EAGER,
+    ATTR_RUST_ARCHETYPE_NATIVE,
 };
 
 // ---
@@ -694,6 +695,10 @@ impl Object {
     pub fn is_eager_rust_archetype(&self) -> bool {
         self.is_archetype() && self.is_attr_set(ATTR_RUST_ARCHETYPE_EAGER)
     }
+
+    pub fn requires_native_rust_archetype(&self) -> bool {
+        self.is_eager_rust_archetype() && self.is_attr_set(ATTR_RUST_ARCHETYPE_NATIVE)
+    }
 }
 
 pub fn is_testing_fqname(fqname: &str) -> bool {
@@ -1426,6 +1431,10 @@ impl Attributes {
     pub fn has(&self, name: impl AsRef<str>) -> bool {
         self.0.contains_key(name.as_ref())
     }
+
+    pub fn remove(&mut self, name: impl AsRef<str>) {
+        self.0.remove(name.as_ref());
+    }
 }
 
 fn filepath_from_declaration_file(

crates/store/re_types/definitions/attributes/rust.fbs

@@ -42,6 +42,14 @@ attribute "attr.rust.new_pub_crate";
 /// object will be generated in `re_viewport/src/blueprint`.
 attribute "attr.rust.override_crate";
 
+/// Whether we should generate an extra Rust object comprised of native Rust types.
+///
+/// The generated object will have the name of the archetype, prefixed by `Native`,
+/// e.g. `NativePoints3D`.
+///
+/// Applies only to archetypes. No-op otherwise.
+attribute "attr.rust.archetype_native";
+
 /// The generated Rust object should be eagerly serialized, i.e. only comprised of Arrow arrays.
 ///
 /// Applies only to archetypes. No-op otherwise.

crates/store/re_types/definitions/rerun/blueprint/archetypes/background.fbs

@@ -5,7 +5,7 @@ namespace rerun.blueprint.archetypes;
 table Background (
     "attr.python.aliases": "datatypes.Rgba32Like, blueprint_components.BackgroundKindLike",
     "attr.rerun.scope": "blueprint",
-    "attr.rust.derive": "Copy"
+    "attr.rust.archetype_eager"
 ) {
     // --- Required ---
 

crates/store/re_types/definitions/rerun/blueprint/archetypes/container_blueprint.fbs

@@ -4,7 +4,9 @@ namespace rerun.blueprint.archetypes;
 
 /// The description of a container.
 table ContainerBlueprint (
-    "attr.rerun.scope": "blueprint"
+    "attr.rerun.scope": "blueprint",
+    "attr.rust.archetype_eager",
+    "attr.rust.archetype_native"
 ) {
     // --- Required ---
 

crates/store/re_types/definitions/rerun/blueprint/archetypes/visual_bounds2d.fbs

@@ -9,7 +9,7 @@ namespace rerun.blueprint.archetypes;
 /// based on the bounding-box of the data or other camera information present in the view.
 table VisualBounds2D (
     "attr.rerun.scope": "blueprint",
-    "attr.rust.derive": "Copy"
+    "attr.rust.archetype_eager": ""
 ) {
     /// Controls the visible range of a 2D view.
     ///