C++ codegen to_arrow_data_type for unions

crates/re_types/source_hash.txt

@@ -1,8 +1,4 @@
 # This is a sha256 hash for all direct and indirect dependencies of this crate's build script.
 # It can be safely removed at anytime to force the build script to run again.
 # Check out build.rs to see how it's computed.
-<<<<<<< HEAD
-70102e74a219445db03e9e5d914687f1756e85b35d93e187dd3346d87b99f806
-=======
-b4f2c2b3b6518fd5c5b8fe2255952c6986895efe58cbc6ae0b56b0384a9ad930
->>>>>>> origin/main
+94f7feacba1a7d75fee69d10294c77e2f4539106e4906f3d58a5db776e11c4be

crates/re_types_builder/src/codegen/cpp/method.rs

@@ -1,7 +1,9 @@
 use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 
-use super::{doc_comment, NEWLINE_TOKEN};
+use crate::Docs;
+
+use super::{doc_comment, quote_docstrings, NEWLINE_TOKEN};
 
 #[derive(Default)]
 pub struct MethodDeclaration {
@@ -49,9 +51,47 @@ impl MethodDeclaration {
     }
 }
 
+#[derive(Default)]
+pub enum MethodDocumentation {
+    #[default]
+    None,
+    String(String),
+    Docs(Docs),
+}
+
+impl From<Docs> for MethodDocumentation {
+    fn from(d: Docs) -> Self {
+        Self::Docs(d)
+    }
+}
+
+impl From<String> for MethodDocumentation {
+    fn from(s: String) -> Self {
+        Self::String(s)
+    }
+}
+
+impl From<&str> for MethodDocumentation {
+    fn from(s: &str) -> Self {
+        Self::String(s.to_owned())
+    }
+}
+
+impl quote::ToTokens for MethodDocumentation {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        match self {
+            Self::None => {}
+            Self::String(s) => {
+                tokens.extend(doc_comment(s));
+            }
+            Self::Docs(docs) => tokens.extend(quote_docstrings(docs)),
+        }
+    }
+}
+
 /// A Cpp struct/class method.
 pub struct Method {
-    pub doc_string: String,
+    pub docs: MethodDocumentation,
     pub declaration: MethodDeclaration,
     pub definition_body: TokenStream,
     pub inline: bool,
@@ -60,7 +100,7 @@ pub struct Method {
 impl Default for Method {
     fn default() -> Self {
         Self {
-            doc_string: String::new(),
+            docs: MethodDocumentation::None,
             declaration: MethodDeclaration::default(),
             definition_body: TokenStream::new(),
             inline: true,
@@ -71,38 +111,33 @@ impl Default for Method {
 impl Method {
     pub fn to_hpp_tokens(&self) -> TokenStream {
         let Self {
-            doc_string,
+            docs,
             declaration,
             definition_body,
             inline: is_inline,
         } = self;
 
-        let quoted_doc = if doc_string.is_empty() {
-            quote! {}
-        } else {
-            doc_comment(doc_string)
-        };
         let declaration = declaration.to_hpp_tokens();
         if *is_inline {
             quote! {
                 #NEWLINE_TOKEN
-                #quoted_doc
+                #docs
                 #declaration {
                     #definition_body
                 }
             }
         } else {
             quote! {
                 #NEWLINE_TOKEN
-                #quoted_doc
+                #docs
                 #declaration;
             }
         }
     }
 
     pub fn to_cpp_tokens(&self, class_or_struct_name: &Ident) -> TokenStream {
         let Self {
-            doc_string: _,
+            docs: _,
             declaration,
             definition_body,
             inline,

crates/re_types_builder/src/codegen/cpp/mod.rs

@@ -217,7 +217,7 @@ impl QuotedObject {
     pub fn new(arrow_registry: &ArrowRegistry, objects: &Objects, obj: &crate::Object) -> Self {
         match obj.specifics {
             crate::ObjectSpecifics::Struct => Self::from_struct(arrow_registry, objects, obj),
-            crate::ObjectSpecifics::Union { .. } => Self::from_union(objects, obj),
+            crate::ObjectSpecifics::Union { .. } => Self::from_union(arrow_registry, objects, obj),
         }
     }
 
@@ -326,7 +326,11 @@ impl QuotedObject {
         Self { hpp, cpp }
     }
 
-    fn from_union(objects: &Objects, obj: &crate::Object) -> QuotedObject {
+    fn from_union(
+        arrow_registry: &ArrowRegistry,
+        objects: &Objects,
+        obj: &crate::Object,
+    ) -> QuotedObject {
         // We implement sum-types as tagged unions;
         // Putting non-POD types in a union requires C++11.
         //
@@ -376,6 +380,9 @@ impl QuotedObject {
         let mut hpp_includes = Includes::default();
         hpp_includes.system.insert("cstdint".to_owned()); // we use `uint32_t` etc everywhere.
         hpp_includes.system.insert("utility".to_owned()); // std::move
+        hpp_includes.system.insert("cstring".to_owned()); // std::memcpy
+        let mut cpp_includes = Includes::default();
+        let mut hpp_declarations = ForwardDecls::default();
 
         let enum_data_declarations = obj
             .fields
@@ -393,44 +400,44 @@ impl QuotedObject {
             })
             .collect_vec();
 
-        let implicit_constructors = if are_types_disjoint(&obj.fields) {
+        let mut methods = Vec::new();
+
+        // Add one static constructor for every field.
+        for obj_field in &obj.fields {
+            methods.push(static_constructor_for_enum_type(
+                objects,
+                &mut hpp_includes,
+                obj_field,
+                &pascal_case_ident,
+                &tag_typename,
+            ));
+        }
+
+        if are_types_disjoint(&obj.fields) {
             // Implicit construct from the different variant types:
-            obj.fields
-                .iter()
-                .map(|obj_field| {
-                    let snake_case_ident =
-                        format_ident!("{}", crate::to_snake_case(&obj_field.name));
-                    let docstring = quote_docstrings(&obj_field.docs);
-                    let param_declaration =
-                        quote_declaration(&mut hpp_includes, obj_field, &snake_case_ident);
-                    quote! {
-                        #docstring
-                        #pascal_case_ident(#param_declaration)
-                        {
-                            *this = #pascal_case_ident::#snake_case_ident(std::move(#snake_case_ident));
-                        }
-                    }
-                })
-                .collect_vec()
+            for obj_field in &obj.fields {
+                let snake_case_ident = format_ident!("{}", crate::to_snake_case(&obj_field.name));
+                let param_declaration =
+                    quote_declaration(&mut hpp_includes, obj_field, &snake_case_ident);
+
+                methods.push(Method {
+                    docs: obj_field.docs.clone().into(),
+                    declaration: MethodDeclaration::constructor(quote!(#pascal_case_ident(#param_declaration))),
+                    definition_body: quote!(*this = #pascal_case_ident::#snake_case_ident(std::move(#snake_case_ident));),
+                    inline: true,
+                });
+            }
         } else {
             // Cannot make implicit constructors, e.g. for
             // `enum Angle { Radians(f32), Degrees(f32) };`
-            vec![]
         };
 
-        let static_constructors = obj
-            .fields
-            .iter()
-            .map(|obj_field| {
-                quote_static_constructor_for_enum_type(
-                    objects,
-                    &mut hpp_includes,
-                    obj_field,
-                    &pascal_case_ident,
-                    &tag_typename,
-                )
-            })
-            .collect_vec();
+        methods.push(arrow_data_type_method(
+            &arrow_registry.get(&obj.fqname),
+            &mut hpp_includes,
+            &mut cpp_includes,
+            &mut hpp_declarations,
+        ));
 
         let destructor = if obj.has_default_destructor(objects) {
             // No destructor needed
@@ -495,10 +502,9 @@ impl QuotedObject {
             }
         };
 
-        hpp_includes.system.insert("cstring".to_owned()); // std::memcpy
-
         let swap_comment = comment("This bitwise swap would fail for self-referential types, but we don't have any of those.");
 
+        let hpp_methods = methods.iter().map(|m| m.to_hpp_tokens());
         let hpp = quote! {
             #hpp_includes
 
@@ -550,9 +556,7 @@ impl QuotedObject {
 
                         #destructor
 
-                        #(#static_constructors)*
-
-                        #(#implicit_constructors)*
+                        #(#hpp_methods)*
 
                         // This is useful for easily implementing the move constructor and move assignment operator:
                         void swap(#pascal_case_ident& other) noexcept {
@@ -569,7 +573,16 @@ impl QuotedObject {
             }
         };
 
-        let cpp = quote! {}; // TODO(emilk): add Arrow serialization code here!
+        let cpp_methods = methods.iter().map(|m| m.to_cpp_tokens(&pascal_case_ident));
+        let cpp = quote! {
+            #cpp_includes
+
+            namespace rr {
+                namespace #namespace_ident {
+                    #(#cpp_methods)*
+                }
+            }
+        };
 
         Self { hpp, cpp }
     }
@@ -588,7 +601,7 @@ fn arrow_data_type_method(
     let quoted_datatype = quote_arrow_data_type(datatype, cpp_includes, true);
 
     Method {
-        doc_string: "Returns the arrow data type this type corresponds to.".to_owned(),
+        docs: "Returns the arrow data type this type corresponds to.".into(),
         declaration: MethodDeclaration {
             is_static: true,
             return_type: quote! { std::shared_ptr<arrow::DataType> },
@@ -600,18 +613,23 @@ fn arrow_data_type_method(
 }
 
 /// e.g. `static Angle radians(float radians);` -> `auto angle = Angle::radians(radians);`
-fn quote_static_constructor_for_enum_type(
+fn static_constructor_for_enum_type(
     objects: &Objects,
     hpp_includes: &mut Includes,
     obj_field: &ObjectField,
     pascal_case_ident: &Ident,
     tag_typename: &Ident,
-) -> TokenStream {
+) -> Method {
     let tag_ident = format_ident!("{}", obj_field.name);
     let snake_case_ident = format_ident!("{}", crate::to_snake_case(&obj_field.name));
-    let docstring = quote_docstrings(&obj_field.docs);
+    let docs = obj_field.docs.clone().into();
 
     let param_declaration = quote_declaration(hpp_includes, obj_field, &snake_case_ident);
+    let declaration = MethodDeclaration {
+        is_static: true,
+        return_type: quote!(#pascal_case_ident),
+        name_and_parameters: quote!(#snake_case_ident(#param_declaration)),
+    };
 
     if let Type::Array { elem_type, length } = &obj_field.typ {
         // We need special casing for constructing arrays:
@@ -632,46 +650,49 @@ fn quote_static_constructor_for_enum_type(
 
         let elem_type = quote_element_type(hpp_includes, elem_type);
 
-        quote! {
-            #docstring
-            static #pascal_case_ident #snake_case_ident(#param_declaration)
-            {
+        Method {
+            docs,
+            declaration,
+            definition_body: quote! {
                 typedef #elem_type TypeAlias;
                 #pascal_case_ident self;
                 self._tag = detail::#tag_typename::#tag_ident;
                 for (size_t i = 0; i < #length; i += 1) {
                     #element_assignment
                 }
                 return std::move(self);
-            }
+            },
+            inline: true,
         }
     } else if obj_field.typ.has_default_destructor(objects) {
         // Generate simpler code for simple types:
-        quote! {
-            #docstring
-            static #pascal_case_ident #snake_case_ident(#param_declaration)
-            {
+        Method {
+            docs,
+            declaration,
+            definition_body: quote! {
                 #pascal_case_ident self;
                 self._tag = detail::#tag_typename::#tag_ident;
                 self._data.#snake_case_ident = std::move(#snake_case_ident);
                 return std::move(self);
-            }
+            },
+            inline: true,
         }
     } else {
         // We need to use placement-new since the union is in an uninitialized state here:
         hpp_includes.system.insert("new".to_owned()); // placement-new
         let typedef_declaration =
             quote_declaration(hpp_includes, obj_field, &format_ident!("TypeAlias"));
-        quote! {
-            #docstring
-            static #pascal_case_ident #snake_case_ident(#param_declaration)
-            {
+        Method {
+            docs,
+            declaration,
+            definition_body: quote! {
                 typedef #typedef_declaration;
                 #pascal_case_ident self;
                 self._tag = detail::#tag_typename::#tag_ident;
                 new (&self._data.#snake_case_ident) TypeAlias(std::move(#snake_case_ident));
                 return std::move(self);
-            }
+            },
+            inline: true,
         }
     }
 }
@@ -892,8 +913,7 @@ fn quote_arrow_data_type(
 
         DataType::Extension(fqname, datatype, _metadata) => {
             // If we're not at the top level, we should have already a `to_arrow_datatype` method that we can relay to.
-            // TODO(andreas): Unions don't have `to_arrow_datatype` yet.
-            if is_top_level_type || matches!(datatype.as_ref(), DataType::Union(..)) {
+            if is_top_level_type {
                 // TODO(andreas): We're no`t emitting the actual extension types here yet which is why we're skipping the extension type at top level.
                 // Currently, we wrap only Components in extension types but this is done in `rerun_c`.
                 // In the future we'll add the extension type here to the schema.