Skip to content

Commit

Permalink
implement support for eager-serialized archetypes (enabled on Points3D)
Browse files Browse the repository at this point in the history
  • Loading branch information
teh-cmc committed Jan 9, 2025
1 parent 455f098 commit 95c2ef1
Show file tree
Hide file tree
Showing 284 changed files with 869 additions and 795 deletions.
1 change: 0 additions & 1 deletion Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -6477,7 +6477,6 @@ dependencies = [
"re_arrow2",
"re_build_tools",
"re_byte_size",
"re_error",
"re_format",
"re_log",
"re_log_types",
Expand Down
174 changes: 142 additions & 32 deletions crates/build/re_types_builder/src/codegen/rust/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,11 +181,11 @@ fn generate_object_file(

code.push_str("\n\n");

code.push_str("use ::re_types_core::external::arrow;\n");
code.push_str("use ::re_types_core::try_serialize_field;\n");
code.push_str("use ::re_types_core::SerializationResult;\n");
code.push_str("use ::re_types_core::{DeserializationResult, DeserializationError};\n");
code.push_str("use ::re_types_core::{ComponentDescriptor, ComponentName};\n");
code.push_str("use ::re_types_core::{ComponentBatch, ComponentBatchCowWithDescriptor};\n");
code.push_str("use ::re_types_core::{ComponentBatch, ComponentBatchCowWithDescriptor, SerializedComponentBatch};\n");

// NOTE: `TokenStream`s discard whitespacing information by definition, so we need to
// inject some of our own when writing to file… while making sure that don't inject
Expand Down Expand Up @@ -324,23 +324,27 @@ fn quote_struct(
quote!(true)
} else {
let quoted_is_pods = obj.fields.iter().map(|obj_field| {
let quoted_field_type = quote_field_type_from_object_field(obj_field);
let quoted_field_type = quote_field_type_from_object_field(obj, obj_field);
quote!(<#quoted_field_type>::is_pod())
});
quote!(#(#quoted_is_pods)&&*)
};

let quoted_is_pod = (!obj.is_eager_rust_archetype()).then_some(quote! {
#[inline]
fn is_pod() -> bool {
#is_pod_impl
}
});

quote! {
impl ::re_byte_size::SizeBytes for #name {
#[inline]
fn heap_size_bytes(&self) -> u64 {
#heap_size_bytes_impl
}

#[inline]
fn is_pod() -> bool {
#is_pod_impl
}
#quoted_is_pod
}
}
};
Expand Down Expand Up @@ -397,7 +401,7 @@ fn quote_union(
let name = format_ident!("{}", re_case::to_pascal_case(&obj_field.name));

let quoted_doc = quote_field_docs(reporter, objects, obj_field);
let quoted_type = quote_field_type_from_object_field(obj_field);
let quoted_type = quote_field_type_from_object_field(obj, obj_field);

if obj_field.typ == Type::Unit {
quote! {
Expand Down Expand Up @@ -431,7 +435,7 @@ fn quote_union(
.iter()
.filter(|obj_field| obj_field.typ != Type::Unit)
.map(|obj_field| {
let quoted_field_type = quote_field_type_from_object_field(obj_field);
let quoted_field_type = quote_field_type_from_object_field(obj, obj_field);
quote!(<#quoted_field_type>::is_pod())
})
.collect();
Expand Down Expand Up @@ -655,7 +659,7 @@ impl ObjectFieldTokenizer<'_> {
let Self(reporter, obj, obj_field) = self;
let quoted_docs = quote_field_docs(reporter, objects, obj_field);
let name = format_ident!("{}", &obj_field.name);
let quoted_type = quote_field_type_from_object_field(obj_field);
let quoted_type = quote_field_type_from_object_field(obj, obj_field);

if is_tuple_struct_from_obj(obj) {
quote! {
Expand Down Expand Up @@ -722,13 +726,16 @@ fn quote_field_type_from_typ(typ: &Type, unwrap: bool) -> (TokenStream, bool) {
(quote!(#obj_field_type), unwrapped)
}

fn quote_field_type_from_object_field(obj_field: &ObjectField) -> TokenStream {
let (quoted_type, _) = quote_field_type_from_typ(&obj_field.typ, false);

if obj_field.is_nullable {
quote!(Option<#quoted_type>)
fn quote_field_type_from_object_field(obj: &Object, obj_field: &ObjectField) -> TokenStream {
if obj.is_eager_rust_archetype() {
quote!(Option<SerializedComponentBatch>)
} else {
quoted_type
let (quoted_type, _) = quote_field_type_from_typ(&obj_field.typ, false);
if obj_field.is_nullable {
quote!(Option<#quoted_type>)
} else {
quoted_type
}
}
}

Expand Down Expand Up @@ -1153,7 +1160,8 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
.map(|field| format_ident!("{}", field.name))
.collect::<Vec<_>>();

let all_component_batches = {
// TODO(cmc): This goes away once all archetypes have been made eager.
let all_native_component_batches = {
std::iter::once(quote! {
Some(Self::indicator())
}).chain(obj.fields.iter().map(|obj_field| {
Expand Down Expand Up @@ -1198,8 +1206,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
quote!{ Some(&self.#field_name as &dyn ComponentBatch) }
};

let archetype_field_name = obj_field.snake_case_name();
let descr_fn_name = format_ident!("descriptor_{archetype_field_name}");
let descr_fn_name = format_ident!("descriptor_{field_name}");

quote! {
(#batch).map(|batch| {
Expand All @@ -1213,7 +1220,37 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
}))
};

let all_deserializers = {
let all_eager_component_batches = {
std::iter::once(quote! {
Some(Self::indicator().serialized())
})
.chain(obj.fields.iter().map(|obj_field| {
let field_name = format_ident!("{}", obj_field.name);
quote!(self.#field_name.clone())
}))
};

let as_components_impl = if obj.is_eager_rust_archetype() {
quote! {
#[inline]
fn as_component_batches_v2(&self) -> Vec<SerializedComponentBatch> {
use ::re_types_core::Archetype as _;
[#(#all_eager_component_batches,)*].into_iter().flatten().collect()
}
}
} else {
quote! {
fn as_component_batches(&self) -> Vec<ComponentBatchCowWithDescriptor<'_>> {
re_tracing::profile_function!();

use ::re_types_core::Archetype as _;
[#(#all_native_component_batches,)*].into_iter().flatten().collect()
}
}
};

// TODO(cmc): This goes away once all archetypes have been made eager.
let all_native_deserializers = {
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 @@ -1291,6 +1328,27 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
})
};

let all_eager_deserializers = {
obj.fields.iter().map(|obj_field| {
let field_name = format_ident!("{}", obj_field.name);
let descr_fn_name = format_ident!("descriptor_{field_name}");

let quoted_deser = quote! {
arrays_by_descr
.get(&Self::#descr_fn_name())
.map(|array| SerializedComponentBatch::new(array.clone(), Self::#descr_fn_name()))
};

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

let all_deserializers = if obj.is_eager_rust_archetype() {
quote!(#(#all_eager_deserializers;)*)
} else {
quote!(#(#all_native_deserializers;)*)
};

quote! {
impl #name {
#(#all_descriptor_methods)*
Expand Down Expand Up @@ -1369,7 +1427,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {

let arrays_by_descr: ::nohash_hasher::IntMap<_, _> = arrow_data.into_iter().collect();

#(#all_deserializers;)*
#all_deserializers

Ok(Self {
#(#quoted_field_names,)*
Expand All @@ -1378,12 +1436,7 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
}

impl ::re_types_core::AsComponents for #name {
fn as_component_batches(&self) -> Vec<ComponentBatchCowWithDescriptor<'_>> {
re_tracing::profile_function!();

use ::re_types_core::Archetype as _;
[#(#all_component_batches,)*].into_iter().flatten().collect()
}
#as_components_impl
}

impl ::re_types_core::ArchetypeReflectionMarker for #name { }
Expand Down Expand Up @@ -1428,7 +1481,7 @@ fn quote_from_impl_from_obj(obj: &Object) -> TokenStream {
let quoted_obj_name = format_ident!("{}", obj.name);
let quoted_obj_field_name = format_ident!("{}", obj_field.name);

let quoted_type = quote_field_type_from_object_field(obj_field);
let quoted_type = quote_field_type_from_object_field(obj, obj_field);

let self_field_access = if obj_is_tuple_struct {
quote!(self.0)
Expand Down Expand Up @@ -1580,7 +1633,8 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
}
});

let quoted_required = required.iter().map(|field| {
// TODO(cmc): This goes away once all archetypes have been made eager.
let quoted_native_required = required.iter().map(|field| {
let field_name = format_ident!("{}", field.name);
let (_, unwrapped) = quote_field_type_from_typ(&field.typ, true);
if unwrapped {
Expand All @@ -1591,6 +1645,25 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
}
});

let quoted_eager_required = required.iter().map(|field| {
let field_name = format_ident!("{}", field.name);
let descr_fn_name = format_ident!("descriptor_{field_name}");

let (_, unwrapped) = quote_field_type_from_typ(&field.typ, true);
if unwrapped {
// This was originally a vec/array!
quote!(#field_name: try_serialize_field(Self::#descr_fn_name(), #field_name))
} else {
quote!(#field_name: try_serialize_field(Self::#descr_fn_name(), [#field_name]))
}
});

let quoted_required = if obj.is_eager_rust_archetype() {
quote!(#(#quoted_eager_required,)*)
} else {
quote!(#(#quoted_native_required,)*)
};

let quoted_optional = optional.iter().map(|field| {
let field_name = format_ident!("{}", field.name);
quote!(#field_name: None)
Expand Down Expand Up @@ -1624,15 +1697,16 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
#[inline]
#fn_new_pub fn new(#(#quoted_params,)*) -> Self {
Self {
#(#quoted_required,)*
#quoted_required
#(#quoted_optional,)*
}
}
}
}
};

let with_methods = optional.iter().map(|field| {
// TODO(cmc): This goes away once all archetypes have been made eager.
let native_with_methods = optional.iter().map(|field| {
// fn with_*()
let field_name = format_ident!("{}", field.name);
let method_name = format_ident!("with_{field_name}");
Expand Down Expand Up @@ -1661,11 +1735,47 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
}
});

let eager_with_methods = optional.iter().map(|field| {
// fn with_*()
let field_name = format_ident!("{}", field.name);
let descr_fn_name = format_ident!("descriptor_{field_name}");
let method_name = format_ident!("with_{field_name}");
let (typ, unwrapped) = quote_field_type_from_typ(&field.typ, true);
let docstring = quote_field_docs(reporter, objects, field);

if unwrapped {
// This was originally a vec/array!
quote! {
#docstring
#[inline]
pub fn #method_name(mut self, #field_name: impl IntoIterator<Item = impl Into<#typ>>) -> Self {
self.#field_name = try_serialize_field(Self::#descr_fn_name(), #field_name);
self
}
}
} else {
quote! {
#docstring
#[inline]
pub fn #method_name(mut self, #field_name: impl Into<#typ>) -> Self {
self.#field_name = try_serialize_field(Self::#descr_fn_name(), [#field_name]);
self
}
}
}
});

let with_methods = if obj.is_eager_rust_archetype() {
quote!(#(#eager_with_methods)*)
} else {
quote!(#(#native_with_methods)*)
};

quote! {
impl #name {
#fn_new

#(#with_methods)*
#with_methods
}
}
}
6 changes: 4 additions & 2 deletions crates/build/re_types_builder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -194,15 +194,17 @@ pub const ATTR_RERUN_EXPERIMENTAL: &str = "attr.rerun.experimental";
pub const ATTR_PYTHON_ALIASES: &str = "attr.python.aliases";
pub const ATTR_PYTHON_ARRAY_ALIASES: &str = "attr.python.array_aliases";

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_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";
pub const ATTR_RUST_NEW_PUB_CRATE: &str = "attr.rust.new_pub_crate";
pub const ATTR_RUST_OVERRIDE_CRATE: &str = "attr.rust.override_crate";
pub const ATTR_RUST_REPR: &str = "attr.rust.repr";
pub const ATTR_RUST_TUPLE_STRUCT: &str = "attr.rust.tuple_struct";
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_DOCS_UNRELEASED: &str = "attr.docs.unreleased";
pub const ATTR_DOCS_CATEGORY: &str = "attr.docs.category";
Expand Down
13 changes: 12 additions & 1 deletion crates/build/re_types_builder/src/objects.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use itertools::Itertools;
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_RERUN_COMPONENT_REQUIRED, ATTR_RERUN_OVERRIDE_TYPE, ATTR_RUST_ARCHETYPE_EAGER,
};

// ---
Expand Down Expand Up @@ -686,6 +686,17 @@ impl Object {
self.kind.plural_snake_case().to_owned()
}
}

pub fn is_archetype(&self) -> bool {
self.kind == ObjectKind::Archetype
}

pub fn is_eager_rust_archetype(&self) -> bool {
self.is_archetype()
&& self
.try_get_attr::<String>(ATTR_RUST_ARCHETYPE_EAGER)
.is_some()
}
}

pub fn is_testing_fqname(fqname: &str) -> bool {
Expand Down
1 change: 0 additions & 1 deletion crates/store/re_types/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,6 @@ testing = []
[dependencies]
# Rerun
re_byte_size.workspace = true
re_error.workspace = true
re_format.workspace = true
re_log.workspace = true
re_log_types.workspace = true
Expand Down
5 changes: 5 additions & 0 deletions crates/store/re_types/definitions/attributes/rust.fbs
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,8 @@ attribute "attr.rust.new_pub_crate";
/// an object of kind `Blueprint` with `attr.rust.override_crate=re_viewport`, the final
/// object will be generated in `re_viewport/src/blueprint`.
attribute "attr.rust.override_crate";

/// The generated Rust object should be eagerly serialized, i.e. only comprised of Arrow arrays.
///
/// Applies only to archetypes. No-op otherwise.
attribute "attr.rust.archetype_eager";
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ namespace rerun.archetypes;
/// \example archetypes/points3d_ui_radius title="Log points with radii given in UI points" image="https://static.rerun.io/point3d_ui_radius/e051a65b4317438bcaea8d0eee016ac9460b5336/1200w.png"
/// \example archetypes/points3d_send_columns title="Send several point clouds with varying point count over time in a single call" image="https://static.rerun.io/points3d_send_columns/633b524a2ee439b0e3afc3f894f4927ce938a3ec/1200w.png" missing="rs"
table Points3D (
"attr.rust.archetype_eager": "",
"attr.rust.derive": "PartialEq",
"attr.docs.category": "Spatial 3D",
"attr.docs.view_types": "Spatial3DView, Spatial2DView: if logged above active projection"
Expand Down
Loading

0 comments on commit 95c2ef1

Please sign in to comment.