Skip to content

Commit b10421a

Browse files
youknowoneyouknowone
authored
Merge pull request #4637 from youknowone/refactor-asnumber
Refactor with AsNumber::number_downcast_exact
2 parents f00e6b6 + 823b528 commit b10421a

4 files changed

Lines changed: 66 additions & 117 deletions

File tree

‎vm/src/builtins/complex.rs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -488,6 +488,10 @@ impl AsNumber for PyComplex {
488488
});
489489
&AS_NUMBER
490490
}
491+
492+
fn clone_exact(zelf: &Py<Self>, vm: &VirtualMachine) -> PyRef<Self> {
493+
vm.ctx.new_complex(zelf.value)
494+
}
491495
}
492496

493497
impl PyComplex {
@@ -502,14 +506,6 @@ impl PyComplex {
502506
Ok(vm.ctx.not_implemented())
503507
}
504508
}
505-
506-
fn number_downcast_exact(number: PyNumber, vm: &VirtualMachine) -> PyRef<PyComplex> {
507-
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
508-
zelf.to_owned()
509-
} else {
510-
vm.ctx.new_complex(Self::number_downcast(number).value)
511-
}
512-
}
513509
}
514510

515511
#[derive(FromArgs)]

‎vm/src/builtins/float.rs‎

Lines changed: 21 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -546,41 +546,34 @@ impl Hashable for PyFloat {
546546
impl AsNumber for PyFloat {
547547
fn as_number() -> &'static PyNumberMethods {
548548
static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
549-
add: atomic_func!(|num, other, vm| PyFloat::number_float_op(
549+
add: atomic_func!(|num, other, vm| PyFloat::number_op(
550550
num,
551551
other,
552-
|a, b| a + b,
552+
|a, b, _vm| a + b,
553553
vm
554554
)),
555-
subtract: atomic_func!(|num, other, vm| PyFloat::number_float_op(
555+
subtract: atomic_func!(|num, other, vm| PyFloat::number_op(
556556
num,
557557
other,
558-
|a, b| a - b,
558+
|a, b, _vm| a - b,
559559
vm
560560
)),
561-
multiply: atomic_func!(|num, other, vm| PyFloat::number_float_op(
561+
multiply: atomic_func!(|num, other, vm| PyFloat::number_op(
562562
num,
563563
other,
564-
|a, b| a * b,
564+
|a, b, _vm| a * b,
565565
vm
566566
)),
567-
remainder: atomic_func!(|num, other, vm| PyFloat::number_general_op(
568-
num, other, inner_mod, vm
569-
)),
570-
divmod: atomic_func!(|num, other, vm| PyFloat::number_general_op(
571-
num,
572-
other,
573-
inner_divmod,
574-
vm
575-
)),
576-
power: atomic_func!(|num, other, vm| PyFloat::number_general_op(
577-
num, other, float_pow, vm
578-
)),
567+
remainder: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, inner_mod, vm)),
568+
divmod: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, inner_divmod, vm)),
569+
power: atomic_func!(|num, other, vm| PyFloat::number_op(num, other, float_pow, vm)),
579570
negative: atomic_func!(|num, vm| {
580571
let value = PyFloat::number_downcast(num).value;
581572
(-value).to_pyresult(vm)
582573
}),
583-
positive: atomic_func!(|num, vm| PyFloat::number_float(num, vm).to_pyresult(vm)),
574+
positive: atomic_func!(
575+
|num, vm| PyFloat::number_downcast_exact(num, vm).to_pyresult(vm)
576+
),
584577
absolute: atomic_func!(|num, vm| {
585578
let value = PyFloat::number_downcast(num).value;
586579
value.abs().to_pyresult(vm)
@@ -590,26 +583,26 @@ impl AsNumber for PyFloat {
590583
let value = PyFloat::number_downcast(num).value;
591584
try_to_bigint(value, vm).map(|x| vm.ctx.new_int(x))
592585
}),
593-
float: atomic_func!(|num, vm| Ok(PyFloat::number_float(num, vm))),
586+
float: atomic_func!(|num, vm| Ok(PyFloat::number_downcast_exact(num, vm))),
594587
floor_divide: atomic_func!(|num, other, vm| {
595-
PyFloat::number_general_op(num, other, inner_floordiv, vm)
588+
PyFloat::number_op(num, other, inner_floordiv, vm)
596589
}),
597590
true_divide: atomic_func!(|num, other, vm| {
598-
PyFloat::number_general_op(num, other, inner_div, vm)
591+
PyFloat::number_op(num, other, inner_div, vm)
599592
}),
600593
..PyNumberMethods::NOT_IMPLEMENTED
601594
});
602595
&AS_NUMBER
603596
}
597+
598+
#[inline]
599+
fn clone_exact(zelf: &Py<Self>, vm: &VirtualMachine) -> PyRef<Self> {
600+
vm.ctx.new_float(zelf.value)
601+
}
604602
}
605603

606604
impl PyFloat {
607-
fn number_general_op<F, R>(
608-
number: PyNumber,
609-
other: &PyObject,
610-
op: F,
611-
vm: &VirtualMachine,
612-
) -> PyResult
605+
fn number_op<F, R>(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
613606
where
614607
F: FnOnce(f64, f64, &VirtualMachine) -> R,
615608
R: ToPyResult,
@@ -620,26 +613,6 @@ impl PyFloat {
620613
Ok(vm.ctx.not_implemented())
621614
}
622615
}
623-
624-
fn number_float_op<F>(
625-
number: PyNumber,
626-
other: &PyObject,
627-
op: F,
628-
vm: &VirtualMachine,
629-
) -> PyResult
630-
where
631-
F: FnOnce(f64, f64) -> f64,
632-
{
633-
Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm)
634-
}
635-
636-
fn number_float(number: PyNumber, vm: &VirtualMachine) -> PyRef<PyFloat> {
637-
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
638-
zelf.to_owned()
639-
} else {
640-
vm.ctx.new_float(Self::number_downcast(number).value)
641-
}
642-
}
643616
}
644617

645618
// Retrieve inner float value:

‎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)]

‎vm/src/types/slot.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,10 +966,24 @@ pub trait AsNumber: PyPayload {
966966
#[pyslot]
967967
fn as_number() -> &'static PyNumberMethods;
968968

969+
fn clone_exact(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
970+
// not all AsNumber requires this implementation.
971+
unimplemented!()
972+
}
973+
969974
#[inline]
970975
fn number_downcast(num: PyNumber) -> &Py<Self> {
971976
unsafe { num.obj.downcast_unchecked_ref() }
972977
}
978+
979+
#[inline]
980+
fn number_downcast_exact(number: PyNumber, vm: &VirtualMachine) -> PyRef<Self> {
981+
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
982+
zelf.to_owned()
983+
} else {
984+
Self::clone_exact(Self::number_downcast(number), vm)
985+
}
986+
}
973987
}
974988

975989
#[pyclass]

0 commit comments

Comments
 (0)