From c76126f951ee582ad4e5dada3ea22f568103752c Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Thu, 24 Sep 2020 03:07:43 +0900 Subject: [PATCH] tp_cmp optimization --- derive/src/pyclass.rs | 7 ------- vm/src/obj/objfunction.rs | 33 ++++++++++++++++++--------------- vm/src/obj/objobject.rs | 26 ++++++++++++++++++-------- vm/src/obj/objtype.rs | 9 ++++++++- vm/src/slots.rs | 25 ++++++++++--------------- vm/src/vm.rs | 10 +++++----- 6 files changed, 59 insertions(+), 51 deletions(-) diff --git a/derive/src/pyclass.rs b/derive/src/pyclass.rs index ae3528937..cd149cddd 100644 --- a/derive/src/pyclass.rs +++ b/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 _)) diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs index 1f5e1765b..4e1be30c9 100644 --- a/vm/src/obj/objfunction.rs +++ b/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, + other: PyObjectRef, + op: PyComparisonOp, + _vm: &VirtualMachine, + ) -> PyResult { + 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 { vm.get_attribute(self.function.clone(), "__module__").ok() } - - #[pymethod(magic)] - fn eq(&self, other: PyObjectRef) -> PyResult { - 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 { - Ok(self.eq(other)?.map(|v| !v)) - } } impl PyValue for PyBoundMethod { diff --git a/vm/src/obj/objobject.rs b/vm/src/obj/objobject.rs index 6b6283cae..f0d33d818 100644 --- a/vm/src/obj/objobject.rs +++ b/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> { + Self::cmp(zelf, other, op, vm).map(Either::B) + } + + #[inline(always)] + fn cmp( + zelf: PyObjectRef, + other: PyObjectRef, + op: PyComparisonOp, + vm: &VirtualMachine, ) -> PyResult { let res = match op { PyComparisonOp::Eq => { @@ -76,7 +86,7 @@ impl PyBaseObject { other: PyObjectRef, vm: &VirtualMachine, ) -> PyResult { - 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 { - 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 { - 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 { - 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 { - 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 { - Self::tp_cmp(zelf, other, PyComparisonOp::Gt, vm) + Self::cmp(zelf, other, PyComparisonOp::Gt, vm) } #[pymethod(magic)] diff --git a/vm/src/obj/objtype.rs b/vm/src/obj/objtype.rs index 3361b2f29..daacc2744 100644 --- a/vm/src/obj/objtype.rs +++ b/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)) + } _ => (), } } diff --git a/vm/src/slots.rs b/vm/src/slots.rs index 2fb21ffa4..4055bb8b0 100644 --- a/vm/src/slots.rs +++ b/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, Option, &VirtualMachine) -> PyResult; pub(crate) type HashFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult; +pub(crate) type CmpFunc = fn( + PyObjectRef, + PyObjectRef, + PyComparisonOp, + &VirtualMachine, +) -> PyResult>; #[derive(Default)] pub struct PyClassSlots { @@ -55,20 +61,9 @@ pub struct PyClassSlots { pub call: AtomicCell>, pub descr_get: AtomicCell>, pub hash: AtomicCell>, - pub cmp: Option, + pub cmp: AtomicCell>, } -type CmpFunc = Box< - py_dyn_fn!( - dyn Fn( - PyObjectRef, - PyObjectRef, - PyComparisonOp, - &VirtualMachine, - ) -> PyResult - ), ->; - 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 { + ) -> PyResult> { 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( diff --git a/vm/src/vm.rs b/vm/src/vm.rs index c4cb06280..b42ec6cf5 100644 --- a/vm/src/vm.rs +++ b/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 };