tp_call optimization

derive/src/pyclass.rs

@@ -342,8 +342,20 @@ where
             let into_func = quote_spanned! {ident.span() =>
                 #transform(Self::#ident)
             };
-            quote! {
-                slots.#slot_ident = Some(#into_func);
+            if slot_name == "call" {
+                quote! {
+                    slots.#slot_ident.store(
+                        Some(
+                            |vm: &::rustpython_vm::VirtualMachine, args: ::rustpython_vm::function::PyFuncArgs| -> ::rustpython_vm::pyobject::PyResult {
+                                ::rustpython_vm::function::IntoPyNativeFunc::call(&Self::#ident, vm, args)
+                            } as _
+                        )
+                    );
+                }
+            } else {
+                quote! {
+                    slots.#slot_ident = Some(#into_func);
+                }
             }
         };
 

vm/src/obj/objtype.rs

@@ -77,6 +77,19 @@ impl PyClassRef {
         std::iter::once(self).chain(self.mro.iter())
     }
 
+    pub(crate) fn first_in_mro<F, R>(&self, f: F) -> Option<R>
+    where
+        F: Fn(&Self) -> Option<R>,
+    {
+        // the hot path will be primitive types which usually hit the result from itself.
+        // try std::intrinsics::likely once it is stablized
+        if let Some(r) = f(self) {
+            Some(r)
+        } else {
+            self.mro.iter().filter_map(|cls| f(cls)).next()
+        }
+    }
+
     pub fn iter_base_chain(&self) -> impl Iterator<Item = &PyClassRef> {
         std::iter::successors(Some(self), |cls| cls.base.as_ref())
     }
@@ -252,7 +265,7 @@ impl PyClass {
                 return Ok(());
             }
         }
-
+        zelf.slots.update_slot_func(attr_name.borrow_value());
         zelf.attributes.write().insert(attr_name.to_string(), value);
         Ok(())
     }
@@ -693,6 +706,11 @@ pub fn new(
     if base.slots.flags.has_feature(PyTpFlags::HAS_DICT) {
         slots.flags |= PyTpFlags::HAS_DICT
     }
+    for slot_name in ["__call__"].iter() {
+        if attrs.contains_key(*slot_name) {
+            slots.update_slot_func(*slot_name);
+        }
+    }
     let new_type = PyRef::new_ref(
         PyClass {
             name: String::from(name),

vm/src/slots.rs

@@ -2,11 +2,12 @@ use std::cmp::Ordering;
 
 use crate::common::cell::PyRwLock;
 use crate::common::hash::PyHash;
-use crate::function::{OptionalArg, PyFuncArgs, PyNativeFunc};
+use crate::function::{IntoPyNativeFunc, OptionalArg, PyFuncArgs, PyNativeFunc};
 use crate::pyobject::{
     IdProtocol, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
 };
 use crate::VirtualMachine;
+use crossbeam_utils::atomic::AtomicCell;
 
 bitflags! {
     pub struct PyTpFlags: u64 {
@@ -44,7 +45,7 @@ pub struct PyClassSlots {
     pub flags: PyTpFlags,
     pub name: PyRwLock<Option<String>>, // tp_name, not class name
     pub new: Option<PyNativeFunc>,
-    pub call: Option<PyNativeFunc>,
+    pub call: AtomicCell<Option<fn(&VirtualMachine, PyFuncArgs) -> PyResult>>,
     pub descr_get: Option<PyDescrGetFunc>,
     pub hash: Option<HashFunc>,
     pub cmp: Option<CmpFunc>,
@@ -69,6 +70,19 @@ impl PyClassSlots {
             ..Default::default()
         }
     }
+
+    #[inline]
+    pub(crate) fn update_slot_func(&self, name: &str) {
+        match name {
+            "__call__" => {
+                self.call
+                    .store(Some(|vm: &VirtualMachine, args: PyFuncArgs| -> PyResult {
+                        IntoPyNativeFunc::call(&call_magic_call, vm, args)
+                    } as _))
+            }
+            _ => (),
+        }
+    }
 }
 
 impl std::fmt::Debug for PyClassSlots {
@@ -84,6 +98,19 @@ pub trait SlotCall: PyValue {
     fn call(zelf: PyRef<Self>, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult;
 }
 
+fn call_magic_call(zelf: PyObjectRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult {
+    use crate::obj::objstr::PyString;
+    use crate::pyobject::IntoPyRef;
+
+    let magic_call = vm.generic_getattribute(zelf, PyString::from("__call__").into_pyref(vm))?;
+    vm.invoke(&magic_call, args)
+
+    // use crate::pyobject::TypeProtocol;
+    // let magic_call = zelf.get_class_attr("__call__").unwrap();
+    // args.insert( zelf);
+    // vm.invoke(&magic_call, args)
+}
+
 pub type PyDescrGetFunc = Box<
     dyn Fn(&VirtualMachine, PyObjectRef, Option<PyObjectRef>, OptionalArg<PyObjectRef>) -> PyResult
         + Send

vm/src/vm.rs

@@ -890,20 +890,21 @@ impl VirtualMachine {
 
     fn _invoke(&self, callable: &PyObjectRef, args: PyFuncArgs) -> PyResult {
         vm_trace!("Invoke: {:?} {:?}", callable, args);
-        if let Some(slot_call) = callable.lease_class().slots.call.as_ref() {
-            self.trace_event(TraceEvent::Call)?;
-            let args = args.insert(callable.clone());
-            let result = slot_call(self, args);
-            self.trace_event(TraceEvent::Return)?;
-            result
-        } else if callable.lease_class().has_attr("__call__") {
-            self.call_method(&callable, "__call__", args)
-        } else {
+        let invoked = callable.class().first_in_mro(|cls| {
+            cls.slots.call.load().map(|slot_call| {
+                self.trace_event(TraceEvent::Call)?;
+                let args = args.insert(callable.clone());
+                let result = slot_call(self, args);
+                self.trace_event(TraceEvent::Return)?;
+                result
+            })
+        });
+        invoked.unwrap_or_else(|| {
             Err(self.new_type_error(format!(
                 "'{}' object is not callable",
                 callable.lease_class().name
             )))
-        }
+        })
     }
 
     #[inline]
@@ -1125,8 +1126,9 @@ impl VirtualMachine {
     }
 
     pub fn is_callable(&self, obj: &PyObjectRef) -> bool {
-        let class = obj.lease_class();
-        class.slots.call.is_some() || class.has_attr("__call__")
+        obj.class()
+            .first_in_mro(|cls| cls.slots.call.load())
+            .is_some()
     }
 
     #[inline]