diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs index a450a07fbe..e98fc5f49b 100644 --- a/vm/src/builtins/complex.rs +++ b/vm/src/builtins/complex.rs @@ -488,6 +488,10 @@ impl AsNumber for PyComplex { }); &AS_NUMBER } + + fn clone_exact(zelf: &Py, vm: &VirtualMachine) -> PyRef { + vm.ctx.new_complex(zelf.value) + } } impl PyComplex { @@ -502,14 +506,6 @@ impl PyComplex { Ok(vm.ctx.not_implemented()) } } - - fn number_downcast_exact(number: PyNumber, vm: &VirtualMachine) -> PyRef { - if let Some(zelf) = number.obj.downcast_ref_if_exact::(vm) { - zelf.to_owned() - } else { - vm.ctx.new_complex(Self::number_downcast(number).value) - } - } } #[derive(FromArgs)] diff --git a/vm/src/builtins/float.rs b/vm/src/builtins/float.rs index 6cb8fb32df..38e1521393 100644 --- a/vm/src/builtins/float.rs +++ b/vm/src/builtins/float.rs @@ -546,41 +546,34 @@ impl Hashable for PyFloat { impl AsNumber for PyFloat { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: Lazy = Lazy::new(|| PyNumberMethods { - add: atomic_func!(|num, other, vm| PyFloat::number_float_op( + add: atomic_func!(|num, other, vm| PyFloat::number_op( num, other, - |a, b| a + b, + |a, b, _vm| a + b, vm )), - subtract: atomic_func!(|num, other, vm| PyFloat::number_float_op( + subtract: atomic_func!(|num, other, vm| PyFloat::number_op( num, other, - |a, b| a - b, + |a, b, _vm| a - b, vm )), - multiply: atomic_func!(|num, other, vm| PyFloat::number_float_op( + multiply: atomic_func!(|num, other, vm| PyFloat::number_op( num, other, - |a, b| a * b, + |a, b, _vm| a * b, vm )), - remainder: atomic_func!(|num, other, vm| PyFloat::number_general_op( - num, other, inner_mod, vm - )), - divmod: atomic_func!(|num, other, vm| PyFloat::number_general_op( - num, - other, - inner_divmod, - vm - )), - power: atomic_func!(|num, other, vm| PyFloat::number_general_op( - num, other, float_pow, vm - )), + remainder: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, inner_mod, vm)), + divmod: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, inner_divmod, vm)), + power: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, float_pow, vm)), negative: atomic_func!(|num, vm| { let value = PyFloat::number_downcast(num).value; (-value).to_pyresult(vm) }), - positive: atomic_func!(|num, vm| PyFloat::number_float(num, vm).to_pyresult(vm)), + positive: atomic_func!( + |num, vm| PyFloat::number_downcast_exact(num, vm).to_pyresult(vm) + ), absolute: atomic_func!(|num, vm| { let value = PyFloat::number_downcast(num).value; value.abs().to_pyresult(vm) @@ -590,26 +583,26 @@ impl AsNumber for PyFloat { let value = PyFloat::number_downcast(num).value; try_to_bigint(value, vm).map(|x| vm.ctx.new_int(x)) }), - float: atomic_func!(|num, vm| Ok(PyFloat::number_float(num, vm))), + float: atomic_func!(|num, vm| Ok(PyFloat::number_downcast_exact(num, vm))), floor_divide: atomic_func!(|num, other, vm| { - PyFloat::number_general_op(num, other, inner_floordiv, vm) + PyFloat::number_op(num, other, inner_floordiv, vm) }), true_divide: atomic_func!(|num, other, vm| { - PyFloat::number_general_op(num, other, inner_div, vm) + PyFloat::number_op(num, other, inner_div, vm) }), ..PyNumberMethods::NOT_IMPLEMENTED }); &AS_NUMBER } + + #[inline] + fn clone_exact(zelf: &Py, vm: &VirtualMachine) -> PyRef { + vm.ctx.new_float(zelf.value) + } } impl PyFloat { - fn number_general_op( - number: PyNumber, - other: &PyObject, - op: F, - vm: &VirtualMachine, - ) -> PyResult + fn number_op(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult where F: FnOnce(f64, f64, &VirtualMachine) -> R, R: ToPyResult, @@ -620,26 +613,6 @@ impl PyFloat { Ok(vm.ctx.not_implemented()) } } - - fn number_float_op( - number: PyNumber, - other: &PyObject, - op: F, - vm: &VirtualMachine, - ) -> PyResult - where - F: FnOnce(f64, f64) -> f64, - { - Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm) - } - - fn number_float(number: PyNumber, vm: &VirtualMachine) -> PyRef { - if let Some(zelf) = number.obj.downcast_ref_if_exact::(vm) { - zelf.to_owned() - } else { - vm.ctx.new_float(Self::number_downcast(number).value) - } - } } // Retrieve inner float value: diff --git a/vm/src/builtins/int.rs b/vm/src/builtins/int.rs index 0443c00eeb..7e9501f693 100644 --- a/vm/src/builtins/int.rs +++ b/vm/src/builtins/int.rs @@ -730,35 +730,26 @@ impl Hashable for PyInt { impl AsNumber for PyInt { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: Lazy = Lazy::new(|| PyNumberMethods { - add: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a + b, vm)), - subtract: atomic_func!(|num, other, vm| PyInt::number_int_op( + add: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a + b, vm)), + subtract: atomic_func!(|num, other, vm| PyInt::number_op( num, other, - |a, b| a - b, + |a, b, _vm| a - b, vm )), - multiply: atomic_func!(|num, other, vm| PyInt::number_int_op( + multiply: atomic_func!(|num, other, vm| PyInt::number_op( num, other, - |a, b| a * b, + |a, b, _vm| a * b, vm )), - remainder: atomic_func!(|num, other, vm| PyInt::number_general_op( - num, other, inner_mod, vm - )), - divmod: atomic_func!(|num, other, vm| PyInt::number_general_op( - num, - other, - inner_divmod, - vm - )), - power: atomic_func!(|num, other, vm| PyInt::number_general_op( - num, other, inner_pow, vm - )), + remainder: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_mod, vm)), + divmod: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_divmod, vm)), + power: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_pow, vm)), negative: atomic_func!(|num, vm| (&PyInt::number_downcast(num).value) .neg() .to_pyresult(vm)), - positive: atomic_func!(|num, vm| Ok(PyInt::number_int(num, vm).into())), + positive: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm).into())), absolute: atomic_func!(|num, vm| PyInt::number_downcast(num) .value .abs() @@ -767,71 +758,46 @@ impl AsNumber for PyInt { invert: atomic_func!(|num, vm| (&PyInt::number_downcast(num).value) .not() .to_pyresult(vm)), - lshift: atomic_func!(|num, other, vm| PyInt::number_general_op( - num, - other, - inner_lshift, - vm - )), - rshift: atomic_func!(|num, other, vm| PyInt::number_general_op( - num, - other, - inner_rshift, - vm - )), - and: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a & b, vm)), - xor: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a ^ b, vm)), - or: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a | b, vm)), - int: atomic_func!(|num, other| Ok(PyInt::number_int(num, other))), + lshift: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_lshift, vm)), + rshift: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_rshift, vm)), + and: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a & b, vm)), + xor: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a ^ b, vm)), + or: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a | b, vm)), + int: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm))), float: atomic_func!(|num, vm| { let zelf = PyInt::number_downcast(num); try_to_float(&zelf.value, vm).map(|x| vm.ctx.new_float(x)) }), floor_divide: atomic_func!(|num, other, vm| { - PyInt::number_general_op(num, other, inner_floordiv, vm) + PyInt::number_op(num, other, inner_floordiv, vm) }), true_divide: atomic_func!(|num, other, vm| { - PyInt::number_general_op(num, other, inner_truediv, vm) + PyInt::number_op(num, other, inner_truediv, vm) }), - index: atomic_func!(|num, vm| Ok(PyInt::number_int(num, vm))), + index: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm))), ..PyNumberMethods::NOT_IMPLEMENTED }); &AS_NUMBER } + + #[inline] + fn clone_exact(zelf: &Py, vm: &VirtualMachine) -> PyRef { + vm.ctx.new_bigint(&zelf.value) + } } impl PyInt { - fn number_general_op( - number: PyNumber, - other: &PyObject, - op: F, - vm: &VirtualMachine, - ) -> PyResult + fn number_op(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult where - F: FnOnce(&BigInt, &BigInt, &VirtualMachine) -> PyResult, + F: FnOnce(&BigInt, &BigInt, &VirtualMachine) -> R, + R: ToPyResult, { if let (Some(a), Some(b)) = (number.obj.payload::(), other.payload::()) { - op(&a.value, &b.value, vm) + op(&a.value, &b.value, vm).to_pyresult(vm) } else { Ok(vm.ctx.not_implemented()) } } - - fn number_int_op(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult - where - F: FnOnce(&BigInt, &BigInt) -> BigInt, - { - Self::number_general_op(number, other, |a, b, _vm| op(a, b).to_pyresult(vm), vm) - } - - fn number_int(number: PyNumber, vm: &VirtualMachine) -> PyIntRef { - if let Some(zelf) = number.obj.downcast_ref_if_exact::(vm) { - zelf.to_owned() - } else { - let zelf = Self::number_downcast(number); - vm.ctx.new_int(zelf.value.clone()) - } - } } #[derive(FromArgs)] diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs index 37c7a23ac8..4369329f4d 100644 --- a/vm/src/types/slot.rs +++ b/vm/src/types/slot.rs @@ -966,10 +966,24 @@ pub trait AsNumber: PyPayload { #[pyslot] fn as_number() -> &'static PyNumberMethods; + fn clone_exact(_zelf: &Py, _vm: &VirtualMachine) -> PyRef { + // not all AsNumber requires this implementation. + unimplemented!() + } + #[inline] fn number_downcast(num: PyNumber) -> &Py { unsafe { num.obj.downcast_unchecked_ref() } } + + #[inline] + fn number_downcast_exact(number: PyNumber, vm: &VirtualMachine) -> PyRef { + if let Some(zelf) = number.obj.downcast_ref_if_exact::(vm) { + zelf.to_owned() + } else { + Self::clone_exact(Self::number_downcast(number), vm) + } + } } #[pyclass]