Skip to content

Commit

Permalink
Add archetype_native attribute and make ContainerBlueprint eager …
Browse files Browse the repository at this point in the history
…+ partial (#8666)

Introduces a new attribute that, when paired with an eager archetype,
will also generate a native companion type with back-and-forth
conversion methods.

This is especially useful for blueprint archetypes (although there are
definitely non-blueprint examples too), which are heavily used all
across the viewer, and would be very painful to use otherwise.

```rust
/// 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 *eager* archetypes. No-op otherwise.
attribute "attr.rust.archetype_native";
```

With this, we now should have all the tools required to port every
remaining archetype (both blueprint and data).

* Part of #8650
  • Loading branch information
teh-cmc authored Jan 14, 2025
1 parent 209547c commit dd04366
Show file tree
Hide file tree
Showing 13 changed files with 554 additions and 230 deletions.
184 changes: 167 additions & 17 deletions crates/build/re_types_builder/src/codegen/rust/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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()) {
Expand All @@ -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);
Expand All @@ -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}");
Expand Down Expand Up @@ -309,10 +335,12 @@ fn quote_struct(
} else {
quote! { pub struct #name { #(#quoted_fields,)* }}
};
let quoted_struct_native = quote_struct_native(reporter, objects, obj);

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);

Expand Down Expand Up @@ -368,6 +396,8 @@ fn quote_struct(
#quoted_deprecation_notice
#quoted_struct

#quoted_struct_native

#quoted_trait_impls

#quoted_from_impl
Expand All @@ -380,6 +410,99 @@ fn quote_struct(
tokens
}

/// Certain eager archetypes might require the generation of an associated native archetype, as
/// the internal viewer code heavily relies on it.
fn quote_struct_native(
reporter: &Reporter,
objects: &Objects,
obj: &Object,
) -> Option<TokenStream> {
assert!(obj.is_struct());

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, .. } = obj;

let name = format_ident!("{name}");

let derive_only = obj.is_attr_set(ATTR_RUST_DERIVE_ONLY);
let quoted_derive_clone_debug = if derive_only {
quote!()
} else {
quote_derive_clone_debug()
};

let is_tuple_struct = is_tuple_struct_from_obj(obj);
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
}
})
}

fn quote_union(
reporter: &Reporter,
arrow_registry: &ArrowRegistry,
Expand Down Expand Up @@ -426,7 +549,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| {
Expand Down Expand Up @@ -574,7 +698,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);
Expand Down Expand Up @@ -843,13 +968,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),
}
Expand Down Expand Up @@ -1049,7 +1175,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 {
Expand Down Expand Up @@ -1238,7 +1364,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);
Expand Down Expand Up @@ -1267,7 +1393,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 {
Expand All @@ -1281,7 +1406,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
Expand All @@ -1291,7 +1416,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)?
Expand All @@ -1304,7 +1429,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)?;

Expand All @@ -1313,7 +1438,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
};

quote!(let #field_name = #quoted_deser;)
})
}).collect_vec()
};

let all_eager_deserializers = {
Expand All @@ -1334,9 +1459,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)*
Expand Down Expand Up @@ -1414,7 +1563,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 {
Expand All @@ -1423,6 +1571,8 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
}
}

#from_arrow_components_native

impl ::re_types_core::AsComponents for #name {
#as_components_impl
}
Expand Down Expand Up @@ -1474,7 +1624,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 {
Expand Down
2 changes: 2 additions & 0 deletions crates/build/re_types_builder/src/codegen/rust/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions crates/build/re_types_builder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
9 changes: 9 additions & 0 deletions crates/build/re_types_builder/src/objects.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

// ---
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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(
Expand Down
Loading

0 comments on commit dd04366

Please sign in to comment.