Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@teh-cmc
teh-cmc committed Jan 9, 2025
1 parent 455f098 commit 38e788c
Show file tree
Hide file tree
Showing 281 changed files with 909 additions and 783 deletions.
201 changes: 169 additions & 32 deletions crates/build/re_types_builder/src/codegen/rust/api.rs
View file
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,20 @@ 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() {
if obj_field.is_nullable {
quote!(Option<SerializedComponentBatch>)
} else {
quote!(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 +1164,9 @@ 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.
// TODO: or does it? what about the native flag?
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 +1211,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 +1225,44 @@ 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);
let is_nullable = obj_field.is_nullable;

if is_nullable {
quote!(self.#field_name.clone())
} else {
quote!(Some(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.
// TODO: or does it? what about the native flag?
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 +1340,40 @@ fn quote_trait_impls_for_archetype(obj: &Object) -> TokenStream {
})
};

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

let is_nullable = obj_field.is_nullable;

let quoted_deser = if is_nullable {
quote! {
arrays_by_descr
.get(&Self::#descr_fn_name())
.map(|array| SerializedComponentBatch::new(array.clone(), Self::#descr_fn_name()))
}
} else {
quote! {{
let array = arrays_by_descr
.get(&Self::#descr_fn_name())
.ok_or_else(DeserializationError::missing_data)
.with_context(#obj_field_fqname)?;
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 +1452,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 +1461,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 +1506,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 +1658,9 @@ 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.
// TODO: or does it? what about the native flag?
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 +1671,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).unwrap())
} else {
quote!(#field_name: try_serialize_field(Self::#descr_fn_name(), [#field_name]).unwrap())
}
});

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 +1723,17 @@ 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.
// TODO: or does it? what about the native flag?
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 +1762,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
}
}
}
7 changes: 5 additions & 2 deletions crates/build/re_types_builder/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,15 +194,18 @@ 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_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";
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
Loading

0 comments on commit 38e788c

Please sign in to comment.