tp_cmp optimization

derive/src/pyclass.rs

@@ -341,13 +341,6 @@ where
                 quote! {
                     slots.#slot_ident = Some(#into_func);
                 }
-            } else if slot_name == "cmp" {
-                let into_func = quote_spanned! {ident.span() =>
-                    ::std::boxed::Box::new(Self::#ident)
-                };
-                quote! {
-                    slots.#slot_ident = Some(#into_func);
-                }
             } else {
                 quote! {
                     slots.#slot_ident.store(Some(Self::#ident as _))

vm/src/obj/objfunction.rs

@@ -19,7 +19,7 @@ use crate::pyobject::{
     PyRef, PyResult, PyValue, TypeProtocol,
 };
 use crate::scope::Scope;
-use crate::slots::{Callable, SlotDescriptor};
+use crate::slots::{Callable, Comparable, PyComparisonOp, SlotDescriptor};
 use crate::VirtualMachine;
 use itertools::Itertools;
 #[cfg(feature = "jit")]
@@ -341,7 +341,23 @@ impl Callable for PyBoundMethod {
     }
 }
 
-#[pyimpl(with(Callable), flags(HAS_DICT))]
+impl Comparable for PyBoundMethod {
+    fn cmp(
+        zelf: PyRef<Self>,
+        other: PyObjectRef,
+        op: PyComparisonOp,
+        _vm: &VirtualMachine,
+    ) -> PyResult<PyComparisonValue> {
+        op.eq_only(|| {
+            let other = class_or_notimplemented!(Self, other);
+            Ok(PyComparisonValue::Implemented(
+                zelf.function.is(&other.function) && zelf.object.is(&other.object),
+            ))
+        })
+    }
+}
+
+#[pyimpl(with(Callable, Comparable), flags(HAS_DICT))]
 impl PyBoundMethod {
     pub fn new(object: PyObjectRef, function: PyObjectRef) -> Self {
         PyBoundMethod { object, function }
@@ -377,19 +393,6 @@ impl PyBoundMethod {
     fn module(&self, vm: &VirtualMachine) -> Option<PyObjectRef> {
         vm.get_attribute(self.function.clone(), "__module__").ok()
     }
-
-    #[pymethod(magic)]
-    fn eq(&self, other: PyObjectRef) -> PyResult<PyComparisonValue> {
-        let other = class_or_notimplemented!(Self, other);
-        Ok(PyComparisonValue::Implemented(
-            self.function.is(&other.function) && self.object.is(&other.object),
-        ))
-    }
-
-    #[pymethod(magic)]
-    fn ne(&self, other: PyObjectRef) -> PyResult<PyComparisonValue> {
-        Ok(self.eq(other)?.map(|v| !v))
-    }
 }
 
 impl PyValue for PyBoundMethod {

vm/src/obj/objobject.rs

@@ -7,8 +7,8 @@ use crate::common::hash::PyHash;
 use crate::function::{OptionalArg, PyFuncArgs};
 use crate::obj::objtype::PyClass;
 use crate::pyobject::{
-    BorrowValue, IdProtocol, ItemProtocol, PyAttributes, PyClassImpl, PyComparisonValue, PyContext,
-    PyObject, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol,
+    BorrowValue, Either, IdProtocol, ItemProtocol, PyAttributes, PyClassImpl, PyComparisonValue,
+    PyContext, PyObject, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol,
 };
 use crate::slots::PyComparisonOp;
 use crate::vm::VirtualMachine;
@@ -44,6 +44,16 @@ impl PyBaseObject {
         other: PyObjectRef,
         op: PyComparisonOp,
         vm: &VirtualMachine,
+    ) -> PyResult<Either<PyObjectRef, PyComparisonValue>> {
+        Self::cmp(zelf, other, op, vm).map(Either::B)
+    }
+
+    #[inline(always)]
+    fn cmp(
+        zelf: PyObjectRef,
+        other: PyObjectRef,
+        op: PyComparisonOp,
+        vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
         let res = match op {
             PyComparisonOp::Eq => {
@@ -76,7 +86,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Eq, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Eq, vm)
     }
     #[pymethod(magic)]
     fn ne(
@@ -84,7 +94,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Ne, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Ne, vm)
     }
     #[pymethod(magic)]
     fn lt(
@@ -92,7 +102,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Lt, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Lt, vm)
     }
     #[pymethod(magic)]
     fn le(
@@ -100,7 +110,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Le, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Le, vm)
     }
     #[pymethod(magic)]
     fn ge(
@@ -108,7 +118,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Ge, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Ge, vm)
     }
     #[pymethod(magic)]
     fn gt(
@@ -116,7 +126,7 @@ impl PyBaseObject {
         other: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<PyComparisonValue> {
-        Self::tp_cmp(zelf, other, PyComparisonOp::Gt, vm)
+        Self::cmp(zelf, other, PyComparisonOp::Gt, vm)
     }
 
     #[pymethod(magic)]

vm/src/obj/objtype.rs

@@ -14,7 +14,7 @@ use super::objtuple::PyTuple;
 use super::objweakref::PyWeak;
 use crate::function::{KwArgs, PyFuncArgs};
 use crate::pyobject::{
-    BorrowValue, IdProtocol, PyAttributes, PyClassImpl, PyContext, PyIterable, PyLease,
+    BorrowValue, Either, IdProtocol, PyAttributes, PyClassImpl, PyContext, PyIterable, PyLease,
     PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
 };
 use crate::slots::{self, Callable, PyClassSlots, PyTpFlags};
@@ -157,6 +157,13 @@ impl PyClassRef {
                 } as _;
                 self.slots.del.store(Some(func));
             }
+            "__eq__" | "__ne__" | "__le__" | "__lt__" | "__ge__" | "__gt__" => {
+                let func: slots::CmpFunc = |zelf, other, op, vm| {
+                    let magic = get_class_magic(&zelf, op.method_name());
+                    vm.invoke(&magic, vec![zelf, other]).map(Either::A)
+                } as _;
+                self.slots.cmp.store(Some(func))
+            }
             _ => (),
         }
     }

vm/src/slots.rs

@@ -4,7 +4,7 @@ use crate::common::cell::PyRwLock;
 use crate::common::hash::PyHash;
 use crate::function::{OptionalArg, PyFuncArgs, PyNativeFunc};
 use crate::pyobject::{
-    IdProtocol, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+    Either, IdProtocol, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
 };
 use crate::VirtualMachine;
 use crossbeam_utils::atomic::AtomicCell;
@@ -45,6 +45,12 @@ pub(crate) type DelFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult<()>;
 pub(crate) type DescrGetFunc =
     fn(PyObjectRef, Option<PyObjectRef>, Option<PyObjectRef>, &VirtualMachine) -> PyResult;
 pub(crate) type HashFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult<PyHash>;
+pub(crate) type CmpFunc = fn(
+    PyObjectRef,
+    PyObjectRef,
+    PyComparisonOp,
+    &VirtualMachine,
+) -> PyResult<Either<PyObjectRef, PyComparisonValue>>;
 
 #[derive(Default)]
 pub struct PyClassSlots {
@@ -55,20 +61,9 @@ pub struct PyClassSlots {
     pub call: AtomicCell<Option<GenericMethod>>,
     pub descr_get: AtomicCell<Option<DescrGetFunc>>,
     pub hash: AtomicCell<Option<HashFunc>>,
-    pub cmp: Option<CmpFunc>,
+    pub cmp: AtomicCell<Option<CmpFunc>>,
 }
 
-type CmpFunc = Box<
-    py_dyn_fn!(
-        dyn Fn(
-            PyObjectRef,
-            PyObjectRef,
-            PyComparisonOp,
-            &VirtualMachine,
-        ) -> PyResult<PyComparisonValue>
-    ),
->;
-
 impl PyClassSlots {
     pub fn from_flags(flags: PyTpFlags) -> Self {
         Self {
@@ -224,9 +219,9 @@ pub trait Comparable: PyValue {
         other: PyObjectRef,
         op: PyComparisonOp,
         vm: &VirtualMachine,
-    ) -> PyResult<PyComparisonValue> {
+    ) -> PyResult<Either<PyObjectRef, PyComparisonValue>> {
         let zelf = PyRef::try_from_object(vm, zelf)?;
-        Self::cmp(zelf, other, op, vm)
+        Self::cmp(zelf, other, op, vm).map(Either::B)
     }
 
     fn cmp(

vm/src/vm.rs

@@ -1399,11 +1399,11 @@ impl VirtualMachine {
         // TODO: _Py_EnterRecursiveCall(tstate, " in comparison")
 
         let call_cmp = |obj: PyObjectRef, other, op| {
-            let ret = if let Some(ref cmp) = obj.class().slots.cmp {
-                Some(cmp(obj, other, op, self)?.map(Either::B))
-            } else if let Some(method) = self.get_method(obj, op.method_name()).transpose()? {
-                let ret = self.invoke(&method, vec![other])?;
-                Some(PyArithmaticValue::from_object(self, ret).map(Either::A))
+            let ret = if let Some(ref cmp) = obj.class().first_in_mro(|cls| cls.slots.cmp.load()) {
+                Some(match cmp(obj, other, op, self)? {
+                    Either::A(obj) => PyArithmaticValue::from_object(self, obj).map(Either::A),
+                    Either::B(arithmatic) => arithmatic.map(Either::B),
+                })
             } else {
                 None
             };