Skip to content

Commit ea95777

Browse files
qingshi163youknowone
authored andcommitted
Simplify AsNumber trait
1 parent 1ed18c0 commit ea95777

File tree

5 files changed

+138
-159
lines changed

5 files changed

+138
-159
lines changed

vm/src/builtins/float.rs

+40-41
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use std::borrow::Cow;
2-
31
use super::{
42
try_bigint_to_f64, PyByteArray, PyBytes, PyInt, PyIntRef, PyStr, PyStrRef, PyType, PyTypeRef,
53
};
@@ -153,9 +151,7 @@ impl Constructor for PyFloat {
153151
let float_val = match arg {
154152
OptionalArg::Missing => 0.0,
155153
OptionalArg::Present(val) => {
156-
if cls.is(vm.ctx.types.float_type) && val.class().is(PyFloat::class(vm)) {
157-
unsafe { val.downcast_unchecked::<PyFloat>().value }
158-
} else if let Some(f) = val.try_float_opt(vm)? {
154+
if let Some(f) = val.try_float_opt(vm)? {
159155
f.value
160156
} else {
161157
float_from_string(val, vm)?
@@ -539,13 +535,40 @@ impl Hashable for PyFloat {
539535
}
540536

541537
impl AsNumber for PyFloat {
542-
fn as_number(_zelf: &crate::Py<Self>, _vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> {
543-
Cow::Borrowed(&Self::NUMBER_METHODS)
544-
}
538+
const AS_NUMBER: PyNumberMethods = PyNumberMethods {
539+
add: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a + b, vm)),
540+
subtract: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a - b, vm)),
541+
multiply: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a * b, vm)),
542+
remainder: Some(|number, other, vm| Self::number_general_op(number, other, inner_mod, vm)),
543+
divmod: Some(|number, other, vm| Self::number_general_op(number, other, inner_divmod, vm)),
544+
power: Some(|number, other, vm| Self::number_general_op(number, other, float_pow, vm)),
545+
negative: Some(|number, vm| {
546+
let value = Self::number_downcast(number).value;
547+
(-value).to_pyresult(vm)
548+
}),
549+
positive: Some(|number, vm| Self::number_float(number, vm).to_pyresult(vm)),
550+
absolute: Some(|number, vm| {
551+
let value = Self::number_downcast(number).value;
552+
value.abs().to_pyresult(vm)
553+
}),
554+
boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())),
555+
int: Some(|number, vm| {
556+
let value = Self::number_downcast(number).value;
557+
try_to_bigint(value, vm).map(|x| vm.ctx.new_int(x))
558+
}),
559+
float: Some(|number, vm| Ok(Self::number_float(number, vm))),
560+
floor_divide: Some(|number, other, vm| {
561+
Self::number_general_op(number, other, inner_floordiv, vm)
562+
}),
563+
true_divide: Some(|number, other, vm| {
564+
Self::number_general_op(number, other, inner_div, vm)
565+
}),
566+
..PyNumberMethods::NOT_IMPLEMENTED
567+
};
545568
}
546569

547570
impl PyFloat {
548-
fn np_general_op<F, R>(
571+
fn number_general_op<F, R>(
549572
number: &PyNumber,
550573
other: &PyObject,
551574
op: F,
@@ -562,49 +585,25 @@ impl PyFloat {
562585
}
563586
}
564587

565-
fn np_float_op<F>(number: &PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
588+
fn number_float_op<F>(
589+
number: &PyNumber,
590+
other: &PyObject,
591+
op: F,
592+
vm: &VirtualMachine,
593+
) -> PyResult
566594
where
567595
F: FnOnce(f64, f64) -> f64,
568596
{
569-
Self::np_general_op(number, other, |a, b, _vm| op(a, b), vm)
597+
Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm)
570598
}
571599

572-
fn np_float(number: &PyNumber, vm: &VirtualMachine) -> PyRef<PyFloat> {
600+
fn number_float(number: &PyNumber, vm: &VirtualMachine) -> PyRef<PyFloat> {
573601
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
574602
zelf.to_owned()
575603
} else {
576604
vm.ctx.new_float(Self::number_downcast(number).value)
577605
}
578606
}
579-
580-
const NUMBER_METHODS: PyNumberMethods = PyNumberMethods {
581-
add: Some(|number, other, vm| Self::np_float_op(number, other, |a, b| a + b, vm)),
582-
subtract: Some(|number, other, vm| Self::np_float_op(number, other, |a, b| a - b, vm)),
583-
multiply: Some(|number, other, vm| Self::np_float_op(number, other, |a, b| a * b, vm)),
584-
remainder: Some(|number, other, vm| Self::np_general_op(number, other, inner_mod, vm)),
585-
divmod: Some(|number, other, vm| Self::np_general_op(number, other, inner_divmod, vm)),
586-
power: Some(|number, other, vm| Self::np_general_op(number, other, float_pow, vm)),
587-
negative: Some(|number, vm| {
588-
let value = Self::number_downcast(number).value;
589-
(-value).to_pyresult(vm)
590-
}),
591-
positive: Some(|number, vm| Self::np_float(number, vm).to_pyresult(vm)),
592-
absolute: Some(|number, vm| {
593-
let value = Self::number_downcast(number).value;
594-
value.abs().to_pyresult(vm)
595-
}),
596-
boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())),
597-
int: Some(|number, vm| {
598-
let value = Self::number_downcast(number).value;
599-
try_to_bigint(value, vm).map(|x| vm.ctx.new_int(x))
600-
}),
601-
float: Some(|number, vm| Ok(Self::np_float(number, vm))),
602-
floor_divide: Some(|number, other, vm| {
603-
Self::np_general_op(number, other, inner_floordiv, vm)
604-
}),
605-
true_divide: Some(|number, other, vm| Self::np_general_op(number, other, inner_div, vm)),
606-
..*PyNumberMethods::not_implemented()
607-
};
608607
}
609608

610609
// Retrieve inner float value:

vm/src/builtins/int.rs

+42-50
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ use bstr::ByteSlice;
1818
use num_bigint::{BigInt, BigUint, Sign};
1919
use num_integer::Integer;
2020
use num_traits::{One, Pow, PrimInt, Signed, ToPrimitive, Zero};
21-
use std::fmt;
22-
use std::{borrow::Cow, ops::Neg};
21+
use std::ops::Neg;
22+
use std::{fmt, ops::Not};
2323

2424
/// int(x=0) -> integer
2525
/// int(x, base=10) -> integer
@@ -756,13 +756,46 @@ impl Hashable for PyInt {
756756
}
757757

758758
impl AsNumber for PyInt {
759-
fn as_number(_zelf: &crate::Py<Self>, _vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> {
760-
Cow::Borrowed(&Self::NUMBER_METHODS)
761-
}
759+
const AS_NUMBER: PyNumberMethods = PyNumberMethods {
760+
add: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a + b, vm)),
761+
subtract: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a - b, vm)),
762+
multiply: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a * b, vm)),
763+
remainder: Some(|number, other, vm| Self::number_general_op(number, other, inner_mod, vm)),
764+
divmod: Some(|number, other, vm| Self::number_general_op(number, other, inner_divmod, vm)),
765+
power: Some(|number, other, vm| Self::number_general_op(number, other, inner_pow, vm)),
766+
negative: Some(|number, vm| (&Self::number_downcast(number).value).neg().to_pyresult(vm)),
767+
positive: Some(|number, vm| Ok(Self::number_int(number, vm).into())),
768+
absolute: Some(|number, vm| Self::number_downcast(number).value.abs().to_pyresult(vm)),
769+
boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())),
770+
invert: Some(|number, vm| (&Self::number_downcast(number).value).not().to_pyresult(vm)),
771+
lshift: Some(|number, other, vm| Self::number_general_op(number, other, inner_lshift, vm)),
772+
rshift: Some(|number, other, vm| Self::number_general_op(number, other, inner_rshift, vm)),
773+
and: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a & b, vm)),
774+
xor: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a ^ b, vm)),
775+
or: Some(|number, other, vm| Self::number_int_op(number, other, |a, b| a | b, vm)),
776+
int: Some(|number, other| Ok(Self::number_int(number, other))),
777+
float: Some(|number, vm| {
778+
let zelf = Self::number_downcast(number);
779+
try_to_float(&zelf.value, vm).map(|x| vm.ctx.new_float(x))
780+
}),
781+
floor_divide: Some(|number, other, vm| {
782+
Self::number_general_op(number, other, inner_floordiv, vm)
783+
}),
784+
true_divide: Some(|number, other, vm| {
785+
Self::number_general_op(number, other, inner_truediv, vm)
786+
}),
787+
index: Some(|number, vm| Ok(Self::number_int(number, vm))),
788+
..PyNumberMethods::NOT_IMPLEMENTED
789+
};
762790
}
763791

764792
impl PyInt {
765-
fn np_general_op<F>(number: &PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
793+
fn number_general_op<F>(
794+
number: &PyNumber,
795+
other: &PyObject,
796+
op: F,
797+
vm: &VirtualMachine,
798+
) -> PyResult
766799
where
767800
F: FnOnce(&BigInt, &BigInt, &VirtualMachine) -> PyResult,
768801
{
@@ -773,62 +806,21 @@ impl PyInt {
773806
}
774807
}
775808

776-
fn np_int_op<F>(number: &PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
809+
fn number_int_op<F>(number: &PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
777810
where
778811
F: FnOnce(&BigInt, &BigInt) -> BigInt,
779812
{
780-
Self::np_general_op(number, other, |a, b, _vm| op(a, b).to_pyresult(vm), vm)
813+
Self::number_general_op(number, other, |a, b, _vm| op(a, b).to_pyresult(vm), vm)
781814
}
782815

783-
fn np_int(number: &PyNumber, vm: &VirtualMachine) -> PyIntRef {
816+
fn number_int(number: &PyNumber, vm: &VirtualMachine) -> PyIntRef {
784817
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
785818
zelf.to_owned()
786819
} else {
787820
let zelf = Self::number_downcast(number);
788821
vm.ctx.new_int(zelf.value.clone())
789822
}
790823
}
791-
792-
const NUMBER_METHODS: PyNumberMethods = PyNumberMethods {
793-
add: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a + b, vm)),
794-
subtract: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a - b, vm)),
795-
multiply: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a * b, vm)),
796-
remainder: Some(|number, other, vm| Self::np_general_op(number, other, inner_mod, vm)),
797-
divmod: Some(|number, other, vm| Self::np_general_op(number, other, inner_divmod, vm)),
798-
power: Some(|number, other, vm| Self::np_general_op(number, other, inner_pow, vm)),
799-
negative: Some(|number, vm| {
800-
Self::number_downcast(number)
801-
.value
802-
.clone()
803-
.neg()
804-
.to_pyresult(vm)
805-
}),
806-
positive: Some(|number, vm| Ok(Self::np_int(number, vm).into())),
807-
absolute: Some(|number, vm| Self::number_downcast(number).value.abs().to_pyresult(vm)),
808-
boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())),
809-
invert: Some(|number, vm| {
810-
let value = Self::number_downcast(number).value.clone();
811-
(!value).to_pyresult(vm)
812-
}),
813-
lshift: Some(|number, other, vm| Self::np_general_op(number, other, inner_lshift, vm)),
814-
rshift: Some(|number, other, vm| Self::np_general_op(number, other, inner_rshift, vm)),
815-
and: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a & b, vm)),
816-
xor: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a ^ b, vm)),
817-
or: Some(|number, other, vm| Self::np_int_op(number, other, |a, b| a | b, vm)),
818-
int: Some(|number, other| Ok(Self::np_int(number, other))),
819-
float: Some(|number, vm| {
820-
let zelf = Self::number_downcast(number);
821-
try_to_float(&zelf.value, vm).map(|x| vm.ctx.new_float(x))
822-
}),
823-
floor_divide: Some(|number, other, vm| {
824-
Self::np_general_op(number, other, inner_floordiv, vm)
825-
}),
826-
true_divide: Some(|number, other, vm| {
827-
Self::np_general_op(number, other, inner_truediv, vm)
828-
}),
829-
index: Some(|number, vm| Ok(Self::np_int(number, vm))),
830-
..*PyNumberMethods::not_implemented()
831-
};
832824
}
833825

834826
#[derive(FromArgs)]

0 commit comments

Comments
 (0)