module: Reuse the user's ItemFn

Summary: Instead of constructing our own `ItemFn`, mostly just reuse the one provided by the user, with some adjustments. The next step in this stack is to add generics support to this stuff, and this makes it much easier to preserve the generics

Reviewed By: ndmitchell

Differential Revision: D73726013

fbshipit-source-id: f4e4687da0831cf8404543222913b23f7638f7a0

Author: JakobDegen

starlark_derive/src/module/parse.rs

@@ -56,26 +56,23 @@ impl ModuleKind {
 }
 
 pub(crate) fn parse(mut input: ItemFn) -> syn::Result<StarModule> {
-    let (module_docstring, attrs) = parse_module_attributes(&input)?;
-    let visibility = input.vis;
-    let sig_span = input.sig.span();
-    let name = input.sig.ident;
+    let module_docstring = parse_module_docstring(&input)?;
 
     if input.sig.inputs.len() != 1 {
         return Err(syn::Error::new(
-            sig_span,
+            input.sig.span(),
             "function must have exactly one argument",
         ));
     }
-    let arg = input.sig.inputs.pop().unwrap();
+    let arg = input.sig.inputs.last_mut().unwrap();
     let arg_span = arg.span();
 
-    let (ty, module_kind) = match arg.into_value() {
-        FnArg::Typed(PatType { ty, .. }) if is_mut_globals_builder(&ty) => {
-            (ty, ModuleKind::Globals)
+    let (pat, module_kind) = match arg {
+        FnArg::Typed(PatType { ty, pat, .. }) if is_mut_globals_builder(&ty) => {
+            (pat, ModuleKind::Globals)
         }
-        FnArg::Typed(PatType { ty, .. }) if is_mut_methods_builder(&ty) => {
-            (ty, ModuleKind::Methods)
+        FnArg::Typed(PatType { ty, pat, .. }) if is_mut_methods_builder(&ty) => {
+            (pat, ModuleKind::Methods)
         }
         _ => {
             return Err(syn::Error::new(
@@ -84,19 +81,36 @@ pub(crate) fn parse(mut input: ItemFn) -> syn::Result<StarModule> {
             ));
         }
     };
+    // FIXME(JakobDegen): Pick one form and enforce it
+    let (syn::Pat::Ident(_) | syn::Pat::Wild(_)) = &mut **pat else {
+        return Err(syn::Error::new(pat.span(), "Expected ident"));
+    };
+    // Replace the argument with a known one - the user can't depend on it anyway
+    *pat = Box::new(syn::Pat::Ident(syn::PatIdent {
+        attrs: Default::default(),
+        by_ref: None,
+        mutability: None,
+        ident: syn::Ident::new("globals_builder", pat.span()),
+        subpat: None,
+    }));
+
+    let stmts = std::mem::replace(
+        &mut input.block,
+        Box::new(syn::Block {
+            brace_token: Default::default(),
+            stmts: Vec::new(),
+        }),
+    )
+    .stmts
+    .into_iter()
+    .map(|stmt| parse_stmt(stmt, module_kind))
+    .collect::<syn::Result<_>>()?;
+
     Ok(StarModule {
         module_kind,
-        visibility,
-        globals_builder: *ty,
-        name,
-        attrs,
+        input,
         docstring: module_docstring,
-        stmts: input
-            .block
-            .stmts
-            .into_iter()
-            .map(|stmt| parse_stmt(stmt, module_kind))
-            .collect::<syn::Result<_>>()?,
+        stmts,
     })
 }
 
@@ -127,22 +141,18 @@ fn is_attribute_docstring(x: &Attribute) -> syn::Result<Option<String>> {
 }
 
 /// Return (docstring, other attributes)
-fn parse_module_attributes(input: &ItemFn) -> syn::Result<(Option<String>, Vec<Attribute>)> {
+fn parse_module_docstring(input: &ItemFn) -> syn::Result<Option<String>> {
     let mut doc_attrs = Vec::new();
-    let mut attrs = Vec::new();
     for attr in &input.attrs {
         if let Some(ds) = is_attribute_docstring(attr)? {
             doc_attrs.push(ds);
-        } else {
-            attrs.push(attr.clone());
         }
     }
-    let docs = if doc_attrs.is_empty() {
-        None
+    if doc_attrs.is_empty() {
+        Ok(None)
     } else {
-        Some(doc_attrs.join("\n"))
-    };
-    Ok((docs, attrs))
+        Ok(Some(doc_attrs.join("\n")))
+    }
 }
 
 fn parse_stmt(stmt: Stmt, module_kind: ModuleKind) -> syn::Result<StarStmt> {

starlark_derive/src/module/render.rs

@@ -42,10 +42,7 @@ pub(crate) fn render(x: StarModule) -> syn::Result<TokenStream> {
 
 fn render_impl(x: StarModule) -> syn::Result<syn::ItemFn> {
     let StarModule {
-        name,
-        globals_builder,
-        visibility,
-        attrs,
+        mut input,
         docstring,
         stmts,
         module_kind,
@@ -56,19 +53,31 @@ fn render_impl(x: StarModule) -> syn::Result<syn::ItemFn> {
         .map(render_stmt)
         .collect::<syn::Result<_>>()?;
     let set_docstring = docstring.map(|ds| quote!(globals_builder.set_docstring(#ds);));
-    Ok(syn::parse_quote! {
-        #( #attrs )*
-        #visibility fn #name(globals_builder: #globals_builder) {
-            fn build(globals_builder: #globals_builder) {
+
+    let inner_fn = syn::ItemFn {
+        attrs: Default::default(),
+        vis: syn::Visibility::Inherited,
+        sig: syn::Signature {
+            ident: syn::Ident::new("build", input.sig.ident.span()),
+            ..input.sig.clone()
+        },
+        block: syn::parse_quote! {
+            {
                 #set_docstring
                 #( #stmts )*
                 // Mute warning if stmts is empty.
                 let _ = globals_builder;
             }
+        },
+    };
+    input.block = syn::parse_quote! {
+        {
+            #inner_fn
             static RES: starlark::environment::#statics = starlark::environment::#statics::new();
             RES.populate(build, globals_builder);
         }
-    })
+    };
+    Ok(input)
 }
 
 fn render_stmt(x: StarStmt) -> syn::Result<syn::Stmt> {

starlark_derive/src/module/typ.rs

@@ -22,7 +22,6 @@ use syn::Attribute;
 use syn::Block;
 use syn::Expr;
 use syn::Type;
-use syn::Visibility;
 use syn::spanned::Spanned;
 
 use crate::module::parse::ModuleKind;
@@ -32,15 +31,12 @@ use crate::module::util::unpack_option;
 
 #[derive(Debug)]
 pub(crate) struct StarModule {
-    pub module_kind: ModuleKind,
-    pub visibility: Visibility,
-    // We reuse the users globals_builder to make sure `use` statements etc
-    // make sense
-    pub globals_builder: Type,
-    pub name: Ident,
-    pub attrs: Vec<Attribute>,
-    pub docstring: Option<String>,
-    pub stmts: Vec<StarStmt>,
+    pub(crate) module_kind: ModuleKind,
+    /// The input `ItemFn`, with the body replaced by an empty block, and the one parameter having
+    /// been renamed to `globals_builder`
+    pub(crate) input: syn::ItemFn,
+    pub(crate) docstring: Option<String>,
+    pub(crate) stmts: Vec<StarStmt>,
 }
 
 #[allow(clippy::large_enum_variant)]