module: Hand write virtualization for NativeAttribute

Summary:
So Rust's type system has a, uhh, limitation:

```
fn foo<T>(){}

fn call_static<U: Fn() + 'static>(u: U) {
    u();
}

pub fn check<'x>() {
    // Doesn't work: lifetime may not live long enough
    // call_static(foo::<&'x ()>);
    call_static(foo::<&'x ()> as fn());
}
```

Function ZSTs, when instantiated with some non-static generics, are not `'static`, while function pointers are.

I'm looking to add generics support to `#[starlark_module]`, so think of the `foo` function as being the implementation for the function/attribute that the user wrote.

This causes problems. `NativeAttribute` values need to be `'static`-ish, but they also need access to generic parameters that might be very non-static as a result of being `StarlarkValue`s.

Working around the limitation above would require me to convert to a function pointer - which I could do, but that would introduce a double virtual call. So to avoid introducing the double virtual call, we explicit split out the `&dyn NativeAttr` into a data value and function pointer. That way, we can use the function pointer that we must explicitly create above to do the single virtual function invocation

Hopefully that made sense

Reviewed By: Nero5023

Differential Revision: D73726014

fbshipit-source-id: 4ee8f71cda1e0cc7bd10e60d1f270b766825efa5

Author: JakobDegen

starlark/src/environment/methods.rs

@@ -33,7 +33,6 @@ use crate::values::FrozenValue;
 use crate::values::FrozenValueTyped;
 use crate::values::Heap;
 use crate::values::Value;
-use crate::values::function::NativeAttr;
 use crate::values::function::NativeAttribute;
 use crate::values::function::NativeMeth;
 use crate::values::function::NativeMethod;
@@ -75,7 +74,7 @@ impl Methods {
     /// of a method, this is the `TyCallable`
     pub(crate) fn get_ty(&self, name: &str) -> Option<Ty> {
         match self.members.get_str(name)? {
-            UnboundValue::Attr(attr, _) => Some(attr.typ.dupe()),
+            UnboundValue::Attr(attr) => Some(attr.typ.dupe()),
             UnboundValue::Method(method, _) => Some(method.ty.dupe()),
         }
     }
@@ -171,40 +170,39 @@ impl MethodsBuilder {
     ) {
         // We want to build an attribute, that ignores its self argument, and does no subsequent allocation.
         let value = self.heap.alloc(value);
-        self.set_attribute_fn(
+        self.members.insert(
             name,
-            true,
-            docstring,
-            V::starlark_type_repr(),
-            move |_, _| Ok(value.to_value()),
+            UnboundValue::Attr(self.heap.alloc_simple_typed(NativeAttribute {
+                speculative_exec_safe: true,
+                docstring,
+                typ: V::starlark_type_repr(),
+                data: Some(value),
+                // SAFETY: Set to `Some` immediately above
+                callable: |value, _, _| Ok(unsafe { value.unwrap_unchecked() }.to_value()),
+            })),
         );
     }
 
-    /// Set an attribute. This function is usually called from code
-    /// generated by `starlark_derive` and rarely needs to be called manually.
-    pub fn set_attribute_fn<F>(
+    /// Set an attribute. Only used by `starlark_module` macro
+    #[doc(hidden)]
+    pub fn set_attribute_fn(
         &mut self,
         name: &str,
         speculative_exec_safe: bool,
         docstring: Option<String>,
         typ: Ty,
-        f: F,
-    ) where
-        F: for<'v> Fn(Value<'v>, &'v Heap) -> crate::Result<Value<'v>> + Send + Sync + 'static,
-    {
+        // The first argument is always `None`
+        f: for<'v> fn(Option<FrozenValue>, Value<'v>, &'v Heap) -> crate::Result<Value<'v>>,
+    ) {
         self.members.insert(
             name,
-            UnboundValue::Attr(
-                FrozenValueTyped::new(self.heap.alloc(NativeAttribute {
-                    speculative_exec_safe,
-                    docstring,
-                    typ,
-                }))
-                .unwrap(),
-                FrozenRef::<dyn NativeAttr + 'static>::new(
-                    self.heap.alloc_any_debug_type_name(f).as_ref(),
-                ),
-            ),
+            UnboundValue::Attr(self.heap.alloc_simple_typed(NativeAttribute {
+                speculative_exec_safe,
+                docstring,
+                typ,
+                data: None,
+                callable: f,
+            })),
         );
     }
 

starlark/src/values/types/function.rs

@@ -114,17 +114,6 @@ pub trait NativeMeth: Send + Sync + 'static {
     ) -> crate::Result<Value<'v>>;
 }
 
-/// A native function that can be evaluated.
-pub trait NativeAttr:
-    for<'v> Fn(Value<'v>, &'v Heap) -> crate::Result<Value<'v>> + Send + Sync + 'static
-{
-}
-
-impl<T> NativeAttr for T where
-    T: for<'v> Fn(Value<'v>, &'v Heap) -> crate::Result<Value<'v>> + Send + Sync + 'static
-{
-}
-
 /// Starlark representation of native (Rust) functions.
 ///
 /// Almost always created with [`#[starlark_module]`](macro@crate::starlark_module).
@@ -288,20 +277,23 @@ pub(crate) struct NativeAttribute {
     pub(crate) speculative_exec_safe: bool,
     pub(crate) docstring: Option<String>,
     pub(crate) typ: Ty,
+    /// Essentially a `&dyn Fn(Value, Heap) -> Result<Value>`, but expanded out by hand, and
+    /// with the `Self` hardcoded to always be `Option<FrozenValue>`
+    ///
+    /// We don't use a enum over whether the callable below requires a `data` argument as that
+    /// would introduce an additional branch when calling the attribute
+    pub(crate) data: Option<FrozenValue>,
+    #[allocative(skip)] // "Not general enough"
+    pub(crate) callable:
+        for<'v> fn(Option<FrozenValue>, Value<'v>, &'v Heap) -> crate::Result<Value<'v>>,
 }
 
 starlark_simple_value!(NativeAttribute);
 
 impl NativeAttribute {
     #[inline]
-    pub(crate) fn invoke_method_impl<'v>(
-        function: &dyn NativeAttr,
-        this: Value<'v>,
-        args: &Arguments<'v, '_>,
-        eval: &mut Evaluator<'v, '_, '_>,
-    ) -> crate::Result<Value<'v>> {
-        let method = function(this, eval.heap())?;
-        method.invoke(args, eval)
+    pub(crate) fn invoke<'v>(&self, this: Value<'v>, heap: &'v Heap) -> crate::Result<Value<'v>> {
+        (self.callable)(self.data, this, heap)
     }
 }
 

starlark/src/values/types/unbound.rs

@@ -31,7 +31,6 @@ use crate::values::Heap;
 use crate::values::Value;
 use crate::values::ValueLike;
 use crate::values::function::BoundMethodGen;
-use crate::values::function::NativeAttr;
 use crate::values::function::NativeAttribute;
 use crate::values::function::NativeMeth;
 use crate::values::function::NativeMethod;
@@ -45,10 +44,7 @@ pub(crate) enum UnboundValue {
         FrozenRef<'static, dyn NativeMeth>,
     ),
     /// An attribute with `this` unbound.
-    Attr(
-        FrozenValueTyped<'static, NativeAttribute>,
-        FrozenRef<'static, dyn NativeAttr>,
-    ),
+    Attr(FrozenValueTyped<'static, NativeAttribute>),
 }
 
 impl Debug for UnboundValue {
@@ -62,7 +58,7 @@ impl UnboundValue {
     pub(crate) fn to_frozen_value(&self) -> FrozenValue {
         match self {
             UnboundValue::Method(m, _) => m.to_frozen_value(),
-            UnboundValue::Attr(a, _) => a.to_frozen_value(),
+            UnboundValue::Attr(a) => a.to_frozen_value(),
         }
     }
 
@@ -73,7 +69,7 @@ impl UnboundValue {
             UnboundValue::Method(m, _) => {
                 Ok(heap.alloc_complex(BoundMethodGen::new(this.to_value(), *m)))
             }
-            UnboundValue::Attr(_, a) => a(this, heap),
+            UnboundValue::Attr(a) => a.invoke(this, heap),
         }
     }
 
@@ -90,9 +86,7 @@ impl UnboundValue {
             Some(span),
             |eval| match self {
                 UnboundValue::Method(_, m) => m.invoke(eval, this, args),
-                UnboundValue::Attr(_, a) => {
-                    NativeAttribute::invoke_method_impl(&**a, this, args, eval)
-                }
+                UnboundValue::Attr(a) => a.invoke(this, eval.heap()),
             },
         )
     }

starlark_derive/src/module/render.rs

@@ -144,6 +144,7 @@ fn render_attr(x: StarAttr) -> syn::Stmt {
     let outer: syn::ItemFn = syn::parse_quote! {
         #[allow(non_snake_case)]
         fn #name<'v>(
+            _ignored: std::option::Option<starlark::values::FrozenValue>,
             #this_value: starlark::values::Value<'v>,
             heap: &'v starlark::values::Heap,
         ) -> starlark::Result<starlark::values::Value<'v>> {