Handle visibility in cube

crates/burn-cube-macros/src/codegen_type/base.rs

@@ -10,7 +10,9 @@ struct TypeCodegen {
  name_expand: syn::Ident,
  fields: Vec<syn::Field>,
  generics: GenericsCodegen,
+ vis: syn::Visibility,
 }
+
 impl TypeCodegen {
  pub fn expand_ty(&self) -> proc_macro2::TokenStream {
  let mut fields = quote::quote! {};
@@ -19,17 +21,19 @@ impl TypeCodegen {
  for field in self.fields.iter() {
  let ident = &field.ident;
  let ty = &field.ty;
+ let vis = &field.vis;
 
  fields.extend(quote! {
- #ident: <#ty as CubeType>::ExpandType,
+ #vis #ident: <#ty as CubeType>::ExpandType,
  });
  }
 
  let generics = self.generics.type_definitions();
+ let vis = &self.vis;
 
  quote! {
  #[derive(Clone)]
- struct #name #generics {
+ #vis struct #name #generics {
  #fields
  }
  }
@@ -42,9 +46,10 @@ impl TypeCodegen {
  for field in self.fields.iter() {
  let ident = &field.ident;
  let ty = &field.ty;
+ let vis = &field.vis;
 
  fields.extend(quote! {
- #ident: <#ty as LaunchArg>::RuntimeArg<'a, R>,
+ #vis #ident: <#ty as LaunchArg>::RuntimeArg<'a, R>,
  });
  }
 
@@ -65,9 +70,10 @@ impl TypeCodegen {
  for field in self.fields.iter() {
  let ident = &field.ident;
  let ty = &field.ty;
+ let vis = &field.vis;
 
  args.extend(quote! {
- #ident: <#ty as LaunchArg>::RuntimeArg<'a, R>,
+ #vis #ident: <#ty as LaunchArg>::RuntimeArg<'a, R>,
  });
  fields.extend(quote! {
  #ident,
@@ -76,11 +82,12 @@ impl TypeCodegen {
 
  let generics_impl = self.generics.all_definitions();
  let generics_use = self.generics.all_in_use();
+ let vis = &self.vis;
 
  quote! {
  impl #generics_impl #name #generics_use {
  /// New kernel
- pub fn new(#args) -> Self {
+ #vis fn new(#args) -> Self {
  Self {
  #fields
  }
@@ -137,12 +144,13 @@ impl TypeCodegen {
  for field in self.fields.iter() {
  let ident = &field.ident;
  let ty = &field.ty;
+ let vis = &field.vis;
 
  body_input.extend(quote! {
- #ident: <#ty as LaunchArg>::compile_input(builder, vectorization),
+ #vis #ident: <#ty as LaunchArg>::compile_input(builder, vectorization),
  });
  body_output.extend(quote! {
- #ident: <#ty as LaunchArg>::compile_output(builder, vectorization),
+ #vis #ident: <#ty as LaunchArg>::compile_output(builder, vectorization),
  });
  }
 
@@ -194,9 +202,12 @@ impl TypeCodegen {
 pub(crate) fn generate_cube_type(ast: &syn::DeriveInput, with_launch: bool) -> TokenStream {
  let name = ast.ident.clone();
  let generics = ast.generics.clone();
+ let visibility = ast.vis.clone();
+
  let name_string = name.to_string();
  let name_expand = Ident::new(format!("{}Expand", name_string).as_str(), name.span());
  let name_launch = Ident::new(format!("{}Launch", name_string).as_str(), name.span());
+
  let mut fields = Vec::new();
 
  match &ast.data {
@@ -215,6 +226,7 @@ pub(crate) fn generate_cube_type(ast: &syn::DeriveInput, with_launch: bool) -> T
  name_expand,
  fields,
  generics: GenericsCodegen::new(generics),
+ vis: visibility,
  };
 
  let expand_ty = codegen.expand_ty();

crates/burn-cube-macros/src/lib.rs

@@ -106,7 +106,7 @@ fn codegen_cube(
  func: &syn::ItemFn,
  variable_tracker: &mut VariableTracker,
 ) -> Result<proc_macro2::TokenStream, proc_macro2::TokenStream> {
- let signature = expand_sig(&func.sig, variable_tracker);
+ let signature = expand_sig(&func.sig, &func.vis, variable_tracker);
  let mut body = quote::quote! {};
 
  for statement in func.block.stmts.iter() {
@@ -148,6 +148,7 @@ fn codegen_cube(
 
 fn expand_sig(
  sig: &syn::Signature,
+ visibility: &syn::Visibility,
  variable_tracker: &mut VariableTracker,
 ) -> proc_macro2::TokenStream {
  let mut inputs = quote::quote!();
@@ -188,6 +189,6 @@ fn expand_sig(
 
  quote::quote! {
  /// Expanded Cube function
- pub fn #ident #generics (context: &mut burn_cube::frontend::CubeContext, #inputs) -> #output
+ #visibility fn #ident #generics (context: &mut burn_cube::frontend::CubeContext, #inputs) -> #output
  }
 }