diff --git a/vm/src/builtins/float.rs b/vm/src/builtins/float.rs index cfa2e7e25..d83439b5e 100644 --- a/vm/src/builtins/float.rs +++ b/vm/src/builtins/float.rs @@ -59,11 +59,11 @@ impl From for PyFloat { impl PyObject { pub fn try_float_opt(&self, vm: &VirtualMachine) -> PyResult>> { - PyNumber::from(self).float_opt(vm) + PyNumber::new(self, vm).float_opt(vm) } pub fn try_float(&self, vm: &VirtualMachine) -> PyResult> { - PyNumber::from(self).float(vm) + PyNumber::new(self, vm).float(vm) } } diff --git a/vm/src/builtins/int.rs b/vm/src/builtins/int.rs index 37405831c..9a1ad9cfa 100644 --- a/vm/src/builtins/int.rs +++ b/vm/src/builtins/int.rs @@ -265,7 +265,7 @@ impl Constructor for PyInt { val }; - PyNumber::from(val.as_ref()) + PyNumber::new(val.as_ref(), vm) .int(vm) .map(|x| x.as_bigint().clone()) } diff --git a/vm/src/function/number.rs b/vm/src/function/number.rs index 62922c85b..4e8219c34 100644 --- a/vm/src/function/number.rs +++ b/vm/src/function/number.rs @@ -82,7 +82,7 @@ impl Deref for ArgIntoFloat { impl TryFromObject for ArgIntoFloat { // Equivalent to PyFloat_AsDouble. fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { - let value = PyNumber::from(obj.as_ref()).float(vm)?.to_f64(); + let value = PyNumber::new(obj.as_ref(), vm).float(vm)?.to_f64(); Ok(ArgIntoFloat { value }) } } diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs index 046ba8fae..f2420f709 100644 --- a/vm/src/protocol/number.rs +++ b/vm/src/protocol/number.rs @@ -1,8 +1,5 @@ -use std::borrow::Cow; - use crate::{ builtins::{int, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr}, - common::lock::OnceCell, function::ArgBytesLike, stdlib::warnings, AsObject, PyObject, PyPayload, PyRef, PyResult, TryFromBorrowedObject, VirtualMachine, @@ -100,45 +97,42 @@ impl PyNumberMethods { pub struct PyNumber<'a> { pub obj: &'a PyObject, // some fast path do not need methods, so we do lazy initialize - methods: OnceCell>, + pub methods: Option<&'static PyNumberMethods>, } -impl<'a> From<&'a PyObject> for PyNumber<'a> { - fn from(obj: &'a PyObject) -> Self { +impl<'a> PyNumber<'a> { + pub fn new(obj: &'a PyObject, vm: &VirtualMachine) -> Self { Self { obj, - methods: OnceCell::new(), + methods: Self::find_methods(obj, vm), } } } impl PyNumber<'_> { - pub fn methods(&self, vm: &VirtualMachine) -> &PyNumberMethods { - &*self.methods_cow(vm) + pub fn find_methods(obj: &PyObject, vm: &VirtualMachine) -> Option<&'static PyNumberMethods> { + obj.class() + .mro_find_map(|x| x.slots.as_number.load()) + .map(|f| f(obj, vm)) } - pub fn methods_cow(&self, vm: &VirtualMachine) -> &Cow<'static, PyNumberMethods> { - self.methods.get_or_init(|| { - self.obj - .class() - .mro_find_map(|x| x.slots.as_number.load()) - .map(|f| f(self.obj, vm)) - .unwrap_or_else(|| Cow::Borrowed(&PyNumberMethods::NOT_IMPLEMENTED)) - }) + pub fn methods(&self) -> &'static PyNumberMethods { + self.methods.unwrap_or(&PyNumberMethods::NOT_IMPLEMENTED) } // PyNumber_Check - pub fn check(&self, vm: &VirtualMachine) -> bool { - let methods = self.methods(vm); - methods.int.is_some() - || methods.index.is_some() - || methods.float.is_some() - || self.obj.payload_is::() + pub fn check(obj: &PyObject, vm: &VirtualMachine) -> bool { + Self::find_methods(obj, vm).map_or(false, |methods| { + methods.int.is_some() + || methods.index.is_some() + || methods.float.is_some() + || obj.payload_is::() + }) } // PyIndex_Check - pub fn is_index(&self, vm: &VirtualMachine) -> bool { - self.methods(vm).index.is_some() + pub fn is_index(&self) -> bool { + self.methods().index.is_some() } pub fn int(&self, vm: &VirtualMachine) -> PyResult { @@ -156,7 +150,7 @@ impl PyNumber<'_> { if let Some(i) = self.obj.downcast_ref_if_exact::(vm) { Ok(i.to_owned()) - } else if let Some(f) = self.methods(vm).int { + } else if let Some(f) = self.methods().int { let ret = f(self, vm)?; if !ret.class().is(PyInt::class(vm)) { warnings::warn( @@ -174,7 +168,7 @@ impl PyNumber<'_> { } else { Ok(ret) } - } else if self.methods(vm).index.is_some() { + } else if self.methods().index.is_some() { self.index(vm) } else if let Ok(Ok(f)) = vm.get_special_method(self.obj.to_owned(), identifier!(vm, __trunc__)) @@ -187,7 +181,7 @@ impl PyNumber<'_> { // vm, // )?; let ret = f.invoke((), vm)?; - PyNumber::from(ret.as_ref()).index(vm).map_err(|_| { + PyNumber::new(ret.as_ref(), vm).index(vm).map_err(|_| { vm.new_type_error(format!( "__trunc__ returned non-Integral (type {})", ret.class() @@ -215,7 +209,7 @@ impl PyNumber<'_> { Ok(Some(i.to_owned())) } else if let Some(i) = self.obj.payload::() { Ok(Some(vm.ctx.new_bigint(i.as_bigint()))) - } else if let Some(f) = self.methods(vm).index { + } else if let Some(f) = self.methods().index { let ret = f(self, vm)?; if !ret.class().is(PyInt::class(vm)) { warnings::warn( @@ -250,7 +244,7 @@ impl PyNumber<'_> { pub fn float_opt(&self, vm: &VirtualMachine) -> PyResult>> { if let Some(float) = self.obj.downcast_ref_if_exact::(vm) { Ok(Some(float.to_owned())) - } else if let Some(f) = self.methods(vm).float { + } else if let Some(f) = self.methods().float { let ret = f(self, vm)?; if !ret.class().is(PyFloat::class(vm)) { warnings::warn( @@ -268,7 +262,7 @@ impl PyNumber<'_> { } else { Ok(Some(ret)) } - } else if self.methods(vm).index.is_some() { + } else if self.methods().index.is_some() { let i = self.index(vm)?; let value = int::try_to_float(i.as_bigint(), vm)?; Ok(Some(vm.ctx.new_float(value))) diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs index 38ad64654..ae3099a28 100644 --- a/vm/src/types/slot.rs +++ b/vm/src/types/slot.rs @@ -139,7 +139,7 @@ impl Default for PyTypeFlags { pub(crate) type GenericMethod = fn(&PyObject, FuncArgs, &VirtualMachine) -> PyResult; pub(crate) type AsMappingFunc = fn(&PyObject, &VirtualMachine) -> &'static PyMappingMethods; -pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> Cow<'static, PyNumberMethods>; +pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> &'static PyNumberMethods; pub(crate) type HashFunc = fn(&PyObject, &VirtualMachine) -> PyResult; // CallFunc = GenericMethod pub(crate) type GetattroFunc = fn(&PyObject, PyStrRef, &VirtualMachine) -> PyResult; @@ -340,43 +340,65 @@ fn as_sequence_generic(zelf: &PyObject, vm: &VirtualMachine) -> &'static PySeque static_as_sequence_generic(has_length, has_ass_item) } -fn as_number_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> { - Cow::Owned(PyNumberMethods { - int: then_some_closure!( - zelf.class().has_attr(identifier!(vm, __int__)), - |num, vm| { - let ret = - vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?; - ret.downcast::().map_err(|obj| { - vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class())) - }) - } - ), - float: then_some_closure!( - zelf.class().has_attr(identifier!(vm, __float__)), - |num, vm| { - let ret = - vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?; - ret.downcast::().map_err(|obj| { - vm.new_type_error(format!( - "__float__ returned non-float (type {})", - obj.class() - )) - }) - } - ), - index: then_some_closure!( - zelf.class().has_attr(identifier!(vm, __index__)), - |num, vm| { - let ret = - vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?; - ret.downcast::().map_err(|obj| { - vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class())) - }) - } - ), - ..PyNumberMethods::NOT_IMPLEMENTED - }) +pub(crate) fn static_as_number_generic( + has_int: bool, + has_float: bool, + has_index: bool, +) -> &'static PyNumberMethods { + static METHODS: &[PyNumberMethods] = &[ + new_generic(false, false, false), + new_generic(true, false, false), + new_generic(false, true, false), + new_generic(true, true, false), + new_generic(false, false, true), + new_generic(true, false, true), + new_generic(false, true, true), + new_generic(true, true, true), + ]; + + fn int(num: &PyNumber, vm: &VirtualMachine) -> PyResult> { + let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?; + ret.downcast::().map_err(|obj| { + vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class())) + }) + } + fn float(num: &PyNumber, vm: &VirtualMachine) -> PyResult> { + let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?; + ret.downcast::().map_err(|obj| { + vm.new_type_error(format!( + "__float__ returned non-float (type {})", + obj.class() + )) + }) + } + fn index(num: &PyNumber, vm: &VirtualMachine) -> PyResult> { + let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?; + ret.downcast::().map_err(|obj| { + vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class())) + }) + } + + const fn new_generic(has_int: bool, has_float: bool, has_index: bool) -> PyNumberMethods { + PyNumberMethods { + int: if has_int { Some(int) } else { None }, + float: if has_float { Some(float) } else { None }, + index: if has_index { Some(index) } else { None }, + ..PyNumberMethods::NOT_IMPLEMENTED + } + } + + let key = bool_int(has_int) | (bool_int(has_float) << 1) | (bool_int(has_index) << 2); + + &METHODS[key] +} + +fn as_number_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> &'static PyNumberMethods { + let (has_int, has_float, has_index) = ( + zelf.class().has_attr(identifier!(vm, __int__)), + zelf.class().has_attr(identifier!(vm, __float__)), + zelf.class().has_attr(identifier!(vm, __index__)), + ); + static_as_number_generic(has_int, has_float, has_index) } fn hash_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult { @@ -1041,8 +1063,8 @@ pub trait AsNumber: PyPayload { #[inline] #[pyslot] - fn as_number(_zelf: &PyObject, _vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> { - Cow::Borrowed(&Self::AS_NUMBER) + fn as_number(_zelf: &PyObject, _vm: &VirtualMachine) -> &'static PyNumberMethods { + &Self::AS_NUMBER } fn number_downcast<'a>(number: &'a PyNumber) -> &'a Py { diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs index 1c325595c..f8af60b7a 100644 --- a/vm/src/vm/vm_ops.rs +++ b/vm/src/vm/vm_ops.rs @@ -10,11 +10,13 @@ use crate::{ /// Collection of operators impl VirtualMachine { pub fn to_index_opt(&self, obj: PyObjectRef) -> Option> { - PyNumber::from(obj.as_ref()).index_opt(self).transpose() + PyNumber::new(obj.as_ref(), self) + .index_opt(self) + .transpose() } pub fn to_index(&self, obj: &PyObject) -> PyResult { - PyNumber::from(obj).index(self) + PyNumber::new(obj, self).index(self) } #[inline]