Skip to content

Commit 823b528

Browse files
youknowoneyouknowone
committed
Refactor AsNumber for PyInt
1 parent c17951e commit 823b528

File tree

1 file changed

+27
-61
lines changed

1 file changed

+27
-61
lines changed

‎vm/src/builtins/int.rs

Lines changed: 27 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -730,35 +730,26 @@ impl Hashable for PyInt {
730730
impl AsNumber for PyInt {
731731
fn as_number() -> &'static PyNumberMethods {
732732
static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
733-
add: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a + b, vm)),
734-
subtract: atomic_func!(|num, other, vm| PyInt::number_int_op(
733+
add: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a + b, vm)),
734+
subtract: atomic_func!(|num, other, vm| PyInt::number_op(
735735
num,
736736
other,
737-
|a, b| a - b,
737+
|a, b, _vm| a - b,
738738
vm
739739
)),
740-
multiply: atomic_func!(|num, other, vm| PyInt::number_int_op(
740+
multiply: atomic_func!(|num, other, vm| PyInt::number_op(
741741
num,
742742
other,
743-
|a, b| a * b,
743+
|a, b, _vm| a * b,
744744
vm
745745
)),
746-
remainder: atomic_func!(|num, other, vm| PyInt::number_general_op(
747-
num, other, inner_mod, vm
748-
)),
749-
divmod: atomic_func!(|num, other, vm| PyInt::number_general_op(
750-
num,
751-
other,
752-
inner_divmod,
753-
vm
754-
)),
755-
power: atomic_func!(|num, other, vm| PyInt::number_general_op(
756-
num, other, inner_pow, vm
757-
)),
746+
remainder: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_mod, vm)),
747+
divmod: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_divmod, vm)),
748+
power: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_pow, vm)),
758749
negative: atomic_func!(|num, vm| (&PyInt::number_downcast(num).value)
759750
.neg()
760751
.to_pyresult(vm)),
761-
positive: atomic_func!(|num, vm| Ok(PyInt::number_int(num, vm).into())),
752+
positive: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm).into())),
762753
absolute: atomic_func!(|num, vm| PyInt::number_downcast(num)
763754
.value
764755
.abs()
@@ -767,71 +758,46 @@ impl AsNumber for PyInt {
767758
invert: atomic_func!(|num, vm| (&PyInt::number_downcast(num).value)
768759
.not()
769760
.to_pyresult(vm)),
770-
lshift: atomic_func!(|num, other, vm| PyInt::number_general_op(
771-
num,
772-
other,
773-
inner_lshift,
774-
vm
775-
)),
776-
rshift: atomic_func!(|num, other, vm| PyInt::number_general_op(
777-
num,
778-
other,
779-
inner_rshift,
780-
vm
781-
)),
782-
and: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a & b, vm)),
783-
xor: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a ^ b, vm)),
784-
or: atomic_func!(|num, other, vm| PyInt::number_int_op(num, other, |a, b| a | b, vm)),
785-
int: atomic_func!(|num, other| Ok(PyInt::number_int(num, other))),
761+
lshift: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_lshift, vm)),
762+
rshift: atomic_func!(|num, other, vm| PyInt::number_op(num, other, inner_rshift, vm)),
763+
and: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a & b, vm)),
764+
xor: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a ^ b, vm)),
765+
or: atomic_func!(|num, other, vm| PyInt::number_op(num, other, |a, b, _vm| a | b, vm)),
766+
int: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm))),
786767
float: atomic_func!(|num, vm| {
787768
let zelf = PyInt::number_downcast(num);
788769
try_to_float(&zelf.value, vm).map(|x| vm.ctx.new_float(x))
789770
}),
790771
floor_divide: atomic_func!(|num, other, vm| {
791-
PyInt::number_general_op(num, other, inner_floordiv, vm)
772+
PyInt::number_op(num, other, inner_floordiv, vm)
792773
}),
793774
true_divide: atomic_func!(|num, other, vm| {
794-
PyInt::number_general_op(num, other, inner_truediv, vm)
775+
PyInt::number_op(num, other, inner_truediv, vm)
795776
}),
796-
index: atomic_func!(|num, vm| Ok(PyInt::number_int(num, vm))),
777+
index: atomic_func!(|num, vm| Ok(PyInt::number_downcast_exact(num, vm))),
797778
..PyNumberMethods::NOT_IMPLEMENTED
798779
});
799780
&AS_NUMBER
800781
}
782+
783+
#[inline]
784+
fn clone_exact(zelf: &Py<Self>, vm: &VirtualMachine) -> PyRef<Self> {
785+
vm.ctx.new_bigint(&zelf.value)
786+
}
801787
}
802788

803789
impl PyInt {
804-
fn number_general_op<F>(
805-
number: PyNumber,
806-
other: &PyObject,
807-
op: F,
808-
vm: &VirtualMachine,
809-
) -> PyResult
790+
fn number_op<F, R>(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
810791
where
811-
F: FnOnce(&BigInt, &BigInt, &VirtualMachine) -> PyResult,
792+
F: FnOnce(&BigInt, &BigInt, &VirtualMachine) -> R,
793+
R: ToPyResult,
812794
{
813795
if let (Some(a), Some(b)) = (number.obj.payload::<Self>(), other.payload::<Self>()) {
814-
op(&a.value, &b.value, vm)
796+
op(&a.value, &b.value, vm).to_pyresult(vm)
815797
} else {
816798
Ok(vm.ctx.not_implemented())
817799
}
818800
}
819-
820-
fn number_int_op<F>(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
821-
where
822-
F: FnOnce(&BigInt, &BigInt) -> BigInt,
823-
{
824-
Self::number_general_op(number, other, |a, b, _vm| op(a, b).to_pyresult(vm), vm)
825-
}
826-
827-
fn number_int(number: PyNumber, vm: &VirtualMachine) -> PyIntRef {
828-
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
829-
zelf.to_owned()
830-
} else {
831-
let zelf = Self::number_downcast(number);
832-
vm.ctx.new_int(zelf.value.clone())
833-
}
834-
}
835801
}
836802

837803
#[derive(FromArgs)]

0 commit comments

Comments
 (0)