module: Make generics work for attrs

Summary:
Like it says in the title

Mostly pretty boring, but requires knowing not to pass late bound lifetimes

Reviewed By: Nero5023

Differential Revision: D73726010

fbshipit-source-id: 5323dd03a6388634e7b20df3c94ab9975e781f23

Author: JakobDegen

starlark/src/tests/derive/module/generic.rs

@@ -15,23 +15,63 @@
  * limitations under the License.
  */
 
+use std::marker::PhantomData;
+
 use starlark_derive::starlark_module;
 
 use crate as starlark;
 use crate::assert::Assert;
 use crate::environment::GlobalsBuilder;
+use crate::environment::MethodsBuilder;
+use crate::typing::Ty;
+use crate::values::AllocValue;
+use crate::values::Heap;
+use crate::values::Value;
+use crate::values::none::NoneType;
+use crate::values::type_repr::StarlarkTypeRepr;
 
 #[starlark_module]
-fn generic_builder<T: Default, U>(globals: &mut GlobalsBuilder)
+fn global_builder<T: Default, U>(globals: &mut GlobalsBuilder)
 where
     U: std::fmt::Display + Default,
 {
     const MY_STR: &str = &U::default().to_string();
 }
 
+struct CustomNone<T>(PhantomData<T>);
+
+impl<T> StarlarkTypeRepr for CustomNone<T> {
+    type Canonical = NoneType;
+
+    /// The representation of a type that a user would use verbatim in starlark type annotations
+    fn starlark_type_repr() -> Ty {
+        NoneType::starlark_type_repr()
+    }
+}
+
+impl<'v, T> AllocValue<'v> for CustomNone<T> {
+    fn alloc_value(self, _heap: &'v Heap) -> Value<'v> {
+        Value::new_none()
+    }
+}
+
+#[starlark_module]
+fn method_builder<T: Default, U>(globals: &mut MethodsBuilder)
+where
+    U: std::fmt::Display + Default,
+{
+    // Just check that this compiles
+    #[starlark(attribute)]
+    fn test_attribute(this: u32) -> starlark::Result<CustomNone<T>> {
+        let _u = U::default().to_string();
+        let _t = T::default();
+        Ok(CustomNone(PhantomData))
+    }
+}
+
 #[test]
 fn test_generic_builder() {
     let mut a = Assert::new();
-    a.globals_add(generic_builder::<u8, u8>);
+    a.globals_add(global_builder::<u8, u8>);
     a.eq("\"0\"", "MY_STR");
 }

starlark_derive/src/module/parse.rs

@@ -36,6 +36,7 @@ use syn::spanned::Spanned;
 
 use crate::module::parse::fun::parse_fun;
 use crate::module::typ::StarConst;
+use crate::module::typ::StarGenerics;
 use crate::module::typ::StarModule;
 use crate::module::typ::StarStmt;
 use crate::module::util::is_type_name;
@@ -107,6 +108,7 @@ pub(crate) fn parse(mut input: ItemFn) -> syn::Result<StarModule> {
     .collect::<syn::Result<_>>()?;
 
     Ok(StarModule {
+        generics: StarGenerics::new(input.sig.generics.clone()),
         module_kind,
         input,
         docstring: module_docstring,

starlark_derive/src/module/render.rs

@@ -30,6 +30,7 @@ use crate::module::typ::SpecialParam;
 use crate::module::typ::StarAttr;
 use crate::module::typ::StarConst;
 use crate::module::typ::StarFun;
+use crate::module::typ::StarGenerics;
 use crate::module::typ::StarModule;
 use crate::module::typ::StarStmt;
 use crate::module::util::ident_string;
@@ -44,11 +45,12 @@ fn render_impl(x: StarModule) -> syn::Result<syn::ItemFn> {
         docstring,
         stmts,
         module_kind,
+        generics,
     } = x;
     let statics = format_ident!("{}", module_kind.statics_type_name());
     let stmts: Vec<_> = stmts
         .into_iter()
-        .map(render_stmt)
+        .map(|s| render_stmt(s, &generics))
         .collect::<syn::Result<_>>()?;
     let set_docstring = docstring.map(|ds| quote!(globals_builder.set_docstring(#ds);));
 
@@ -68,22 +70,21 @@ fn render_impl(x: StarModule) -> syn::Result<syn::ItemFn> {
             }
         },
     };
-    let turbo = input.sig.generics.split_for_impl().1;
-    let turbo = turbo.as_turbofish();
+    let turbofish = generics.turbofish();
     input.block = syn::parse_quote! {
         {
             #inner_fn
             static RES: starlark::environment::#statics = starlark::environment::#statics::new();
-            RES.populate(build #turbo, globals_builder);
+            RES.populate(build #turbofish, globals_builder);
         }
     };
     Ok(input)
 }
 
-fn render_stmt(x: StarStmt) -> syn::Result<syn::Stmt> {
+fn render_stmt(x: StarStmt, generics: &StarGenerics) -> syn::Result<syn::Stmt> {
     match x {
         StarStmt::Const(x) => Ok(render_const(x)),
-        StarStmt::Attr(x) => Ok(render_attr(x)),
+        StarStmt::Attr(x) => Ok(render_attr(x, generics)),
         StarStmt::Fun(x) => render_fun(x),
     }
 }
@@ -96,7 +97,7 @@ fn render_const(x: StarConst) -> syn::Stmt {
     }
 }
 
-fn render_attr(x: StarAttr) -> syn::Stmt {
+fn render_attr(x: StarAttr, generics: &StarGenerics) -> syn::Stmt {
     let StarAttr {
         name,
         this,
@@ -114,6 +115,10 @@ fn render_attr(x: StarAttr) -> syn::Stmt {
         None => render_none(),
     };
 
+    let generic_decls = generics.decls();
+    let where_clause = generics.where_clause();
+    let turbofish = generics.turbofish();
+
     let let_heap = if let Some(SpecialParam {
         param: SimpleParam { ident, ty, .. },
     }) = heap
@@ -132,24 +137,28 @@ fn render_attr(x: StarAttr) -> syn::Stmt {
         #( #attrs )*
         #[allow(non_snake_case)] // Starlark doesn't have this convention
         #[allow(unused_variables)]
-        fn #name_inner<'v>(
+        fn #name_inner #generic_decls(
             this: #this_return_type,
             __heap: &'v starlark::values::Heap,
-        ) -> #return_type {
+        ) -> #return_type
+        #where_clause
+        {
             #let_heap
             #body
         }
     };
 
     let outer: syn::ItemFn = syn::parse_quote! {
         #[allow(non_snake_case)]
-        fn #name<'v>(
+        fn #name #generic_decls(
             _ignored: std::option::Option<starlark::values::FrozenValue>,
             #this_value: starlark::values::Value<'v>,
             heap: &'v starlark::values::Heap,
-        ) -> starlark::Result<starlark::values::Value<'v>> {
+        ) -> starlark::Result<starlark::values::Value<'v>>
+        #where_clause
+        {
             #unpack
-            Ok(heap.alloc(#name_inner(this, heap)?))
+            Ok(heap.alloc(#name_inner #turbofish(this, heap)?))
         }
     };
 
@@ -162,8 +171,8 @@ fn render_attr(x: StarAttr) -> syn::Stmt {
                 #name_str,
                 #speculative_exec_safe,
                 #docstring,
-                starlark::values::type_repr::type_repr_from_attr_impl(#name_inner),
-                #name
+                starlark::values::type_repr::type_repr_from_attr_impl(#name_inner #turbofish),
+                #name #turbofish
             );
         }
     }

starlark_derive/src/module/typ.rs

@@ -37,6 +37,47 @@ pub(crate) struct StarModule {
     pub(crate) input: syn::ItemFn,
     pub(crate) docstring: Option<String>,
     pub(crate) stmts: Vec<StarStmt>,
+    pub(crate) generics: StarGenerics,
+}
+
+/// The generics the user provided on the starlark module
+#[derive(Debug)]
+pub(crate) struct StarGenerics {
+    /// The user provided generics
+    ///
+    /// This is what we use to instantiate functions we define with their generics, because the `'v`
+    /// lifetime is late bound and therefore never passed as an explicit type parameter
+    generics: syn::Generics,
+    /// The user provided generics with a `'v` param added in front.
+    ///
+    /// This is what we use to *declare* generics on functions that we define, because all of our
+    /// functions want a `'v` param
+    generics_with_v: syn::Generics,
+}
+
+impl StarGenerics {
+    pub(crate) fn new(g: syn::Generics) -> Self {
+        let mut with_v = g.clone();
+        with_v
+            .params
+            .insert(0, syn::GenericParam::Lifetime(syn::parse_quote! { 'v }));
+        Self {
+            generics: g,
+            generics_with_v: with_v,
+        }
+    }
+
+    pub(crate) fn decls(&self) -> syn::ImplGenerics<'_> {
+        self.generics_with_v.split_for_impl().0
+    }
+
+    pub(crate) fn turbofish(&self) -> syn::Turbofish<'_> {
+        self.generics.split_for_impl().1.as_turbofish()
+    }
+
+    pub(crate) fn where_clause(&self) -> Option<&syn::WhereClause> {
+        self.generics_with_v.split_for_impl().2
+    }
 }
 
 #[allow(clippy::large_enum_variant)]