reintroduce vectorcall optimization with new PyCallArgs trait

PyO3 · Dec 5, 2024 · 0608fa9 · 0608fa9
1 parent 77202aa
commit 0608fa9
Show file tree

Hide file tree

Showing 4 changed files with 401 additions and 29 deletions.
diff --git a/src/call.rs b/src/call.rs
@@ -0,0 +1,120 @@
+//! TODO
+//!
+//!
+
+use crate::ffi_ptr_ext::FfiPtrExt as _;
+use crate::types::{PyAnyMethods as _, PyDict, PyString, PyTuple};
+use crate::{ffi, Borrowed, Bound, IntoPyObjectExt as _, Py, PyAny, PyResult};
+
+pub(crate) mod private {
+    use super::*;
+
+    pub trait Sealed {}
+
+    impl Sealed for () {}
+    impl Sealed for Bound<'_, PyTuple> {}
+    impl Sealed for Py<PyTuple> {}
+
+    pub struct Token;
+}
+
+/// TODO
+pub trait PyCallArgs<'py>: Sized + private::Sealed {
+    #[doc(hidden)]
+    fn call(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Borrowed<'_, 'py, PyDict>,
+        token: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>>;
+
+    #[doc(hidden)]
+    fn call_positional(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        token: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>>;
+
+    #[doc(hidden)]
+    fn call_method_positional(
+        self,
+        object: Borrowed<'_, 'py, PyAny>,
+        method_name: Borrowed<'_, 'py, PyString>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        object
+            .getattr(method_name)
+            .and_then(|method| method.call1(self))
+    }
+}
+
+impl<'py> PyCallArgs<'py> for () {
+    fn call(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Borrowed<'_, 'py, PyDict>,
+        token: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let args = self.into_pyobject_or_pyerr(function.py())?;
+        args.call(function, kwargs, token)
+    }
+
+    fn call_positional(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        token: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let args = self.into_pyobject_or_pyerr(function.py())?;
+        args.call_positional(function, token)
+    }
+}
+
+impl<'py> PyCallArgs<'py> for Bound<'py, PyTuple> {
+    fn call(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Borrowed<'_, '_, PyDict>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            ffi::PyObject_Call(function.as_ptr(), self.as_ptr(), kwargs.as_ptr())
+                .assume_owned_or_err(function.py())
+        }
+    }
+
+    fn call_positional(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            ffi::PyObject_Call(function.as_ptr(), self.as_ptr(), std::ptr::null_mut())
+                .assume_owned_or_err(function.py())
+        }
+    }
+}
+
+impl<'py> PyCallArgs<'py> for Py<PyTuple> {
+    fn call(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        kwargs: Borrowed<'_, '_, PyDict>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            ffi::PyObject_Call(function.as_ptr(), self.as_ptr(), kwargs.as_ptr())
+                .assume_owned_or_err(function.py())
+        }
+    }
+
+    fn call_positional(
+        self,
+        function: Borrowed<'_, 'py, PyAny>,
+        _: private::Token,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        unsafe {
+            ffi::PyObject_Call(function.as_ptr(), self.as_ptr(), std::ptr::null_mut())
+                .assume_owned_or_err(function.py())
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -428,6 +428,7 @@ mod internal_tricks;
 mod internal;
 
 pub mod buffer;
+pub mod call;
 pub mod conversion;
 mod conversions;
 #[cfg(feature = "experimental-async")]

diff --git a/src/types/any.rs b/src/types/any.rs
@@ -1,3 +1,4 @@
+use crate::call::PyCallArgs;
 use crate::class::basic::CompareOp;
 use crate::conversion::{AsPyPointer, FromPyObjectBound, IntoPyObject};
 use crate::err::{DowncastError, DowncastIntoError, PyErr, PyResult};
@@ -10,7 +11,7 @@ use crate::py_result_ext::PyResultExt;
 use crate::type_object::{PyTypeCheck, PyTypeInfo};
 #[cfg(not(any(PyPy, GraalPy)))]
 use crate::types::PySuper;
-use crate::types::{PyDict, PyIterator, PyList, PyString, PyTuple, PyType};
+use crate::types::{PyDict, PyIterator, PyList, PyString, PyType};
 use crate::{err, ffi, Borrowed, BoundObject, IntoPyObjectExt, Python};
 use std::cell::UnsafeCell;
 use std::cmp::Ordering;
@@ -436,7 +437,7 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed {
     /// ```
     fn call<A>(&self, args: A, kwargs: Option<&Bound<'py, PyDict>>) -> PyResult<Bound<'py, PyAny>>
     where
-        A: IntoPyObject<'py, Target = PyTuple>;
+        A: PyCallArgs<'py>;
 
     /// Calls the object without arguments.
     ///
@@ -491,7 +492,7 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed {
     /// ```
     fn call1<A>(&self, args: A) -> PyResult<Bound<'py, PyAny>>
     where
-        A: IntoPyObject<'py, Target = PyTuple>;
+        A: PyCallArgs<'py>;
 
     /// Calls a method on the object.
     ///
@@ -538,7 +539,7 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed {
     ) -> PyResult<Bound<'py, PyAny>>
     where
         N: IntoPyObject<'py, Target = PyString>,
-        A: IntoPyObject<'py, Target = PyTuple>;
+        A: PyCallArgs<'py>;
 
     /// Calls a method on the object without arguments.
     ///
@@ -614,7 +615,7 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed {
     fn call_method1<N, A>(&self, name: N, args: A) -> PyResult<Bound<'py, PyAny>>
     where
         N: IntoPyObject<'py, Target = PyString>,
-        A: IntoPyObject<'py, Target = PyTuple>;
+        A: PyCallArgs<'py>;
 
     /// Returns whether the object is considered to be true.
     ///
@@ -1209,25 +1210,17 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
 
     fn call<A>(&self, args: A, kwargs: Option<&Bound<'py, PyDict>>) -> PyResult<Bound<'py, PyAny>>
     where
-        A: IntoPyObject<'py, Target = PyTuple>,
+        A: PyCallArgs<'py>,
     {
-        fn inner<'py>(
-            any: &Bound<'py, PyAny>,
-            args: Borrowed<'_, 'py, PyTuple>,
-            kwargs: Option<&Bound<'py, PyDict>>,
-        ) -> PyResult<Bound<'py, PyAny>> {
-            unsafe {
-                ffi::PyObject_Call(
-                    any.as_ptr(),
-                    args.as_ptr(),
-                    kwargs.map_or(std::ptr::null_mut(), |dict| dict.as_ptr()),
-                )
-                .assume_owned_or_err(any.py())
-            }
+        if let Some(kwargs) = kwargs {
+            args.call(
+                self.as_borrowed(),
+                kwargs.as_borrowed(),
+                crate::call::private::Token,
+            )
+        } else {
+            args.call_positional(self.as_borrowed(), crate::call::private::Token)
         }
-
-        let py = self.py();
-        inner(self, args.into_pyobject_or_pyerr(py)?.as_borrowed(), kwargs)
     }
 
     #[inline]
@@ -1237,9 +1230,9 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
 
     fn call1<A>(&self, args: A) -> PyResult<Bound<'py, PyAny>>
     where
-        A: IntoPyObject<'py, Target = PyTuple>,
+        A: PyCallArgs<'py>,
     {
-        self.call(args, None)
+        args.call_positional(self.as_borrowed(), crate::call::private::Token)
     }
 
     #[inline]
@@ -1251,10 +1244,14 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
     ) -> PyResult<Bound<'py, PyAny>>
     where
         N: IntoPyObject<'py, Target = PyString>,
-        A: IntoPyObject<'py, Target = PyTuple>,
+        A: PyCallArgs<'py>,
     {
-        self.getattr(name)
-            .and_then(|method| method.call(args, kwargs))
+        if kwargs.is_none() {
+            self.call_method1(name, args)
+        } else {
+            self.getattr(name)
+                .and_then(|method| method.call(args, kwargs))
+        }
     }
 
     #[inline]
@@ -1273,9 +1270,14 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
     fn call_method1<N, A>(&self, name: N, args: A) -> PyResult<Bound<'py, PyAny>>
     where
         N: IntoPyObject<'py, Target = PyString>,
-        A: IntoPyObject<'py, Target = PyTuple>,
+        A: PyCallArgs<'py>,
     {
-        self.call_method(name, args, None)
+        let name = name.into_pyobject_or_pyerr(self.py())?;
+        args.call_method_positional(
+            self.as_borrowed(),
+            name.as_borrowed(),
+            crate::call::private::Token,
+        )
     }
 
     fn is_truthy(&self) -> PyResult<bool> {