Skip to content

Commit 0ebd4e0

Browse files
qingshi163youknowone
authored andcommitted
refactor number methods as atomic function
1 parent c5d6ef1 commit 0ebd4e0

File tree

3 files changed

+200
-177
lines changed

3 files changed

+200
-177
lines changed

vm/src/builtins/float.rs

Lines changed: 43 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use crate::{
1616
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
1717
TryFromBorrowedObject, TryFromObject, VirtualMachine,
1818
};
19+
use crossbeam_utils::atomic::AtomicCell;
1920
use num_bigint::{BigInt, ToBigInt};
2021
use num_complex::Complex64;
2122
use num_rational::Ratio;
@@ -552,34 +553,57 @@ impl Hashable for PyFloat {
552553
}
553554
}
554555

556+
macro_rules! atomic_func {
557+
($x:expr) => {
558+
AtomicCell::new(Some($x))
559+
};
560+
}
561+
555562
impl AsNumber for PyFloat {
556563
const AS_NUMBER: PyNumberMethods = PyNumberMethods {
557-
add: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a + b, vm)),
558-
subtract: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a - b, vm)),
559-
multiply: Some(|number, other, vm| Self::number_float_op(number, other, |a, b| a * b, vm)),
560-
remainder: Some(|number, other, vm| Self::number_general_op(number, other, inner_mod, vm)),
561-
divmod: Some(|number, other, vm| Self::number_general_op(number, other, inner_divmod, vm)),
562-
power: Some(|number, other, vm| Self::number_general_op(number, other, float_pow, vm)),
563-
negative: Some(|number, vm| {
564-
let value = Self::number_downcast(number).value;
564+
add: atomic_func!(|num, other, vm| Self::number_float_op(num, other, |a, b| a + b, vm)),
565+
subtract: atomic_func!(|num, other, vm| Self::number_float_op(
566+
num,
567+
other,
568+
|a, b| a - b,
569+
vm
570+
)),
571+
multiply: atomic_func!(|num, other, vm| Self::number_float_op(
572+
num,
573+
other,
574+
|a, b| a * b,
575+
vm
576+
)),
577+
remainder: atomic_func!(|num, other, vm| Self::number_general_op(
578+
num, other, inner_mod, vm
579+
)),
580+
divmod: atomic_func!(|num, other, vm| Self::number_general_op(
581+
num,
582+
other,
583+
inner_divmod,
584+
vm
585+
)),
586+
power: atomic_func!(|num, other, vm| Self::number_general_op(num, other, float_pow, vm)),
587+
negative: atomic_func!(|num, vm| {
588+
let value = Self::number_downcast(num).value;
565589
(-value).to_pyresult(vm)
566590
}),
567-
positive: Some(|number, vm| Self::number_float(number, vm).to_pyresult(vm)),
568-
absolute: Some(|number, vm| {
569-
let value = Self::number_downcast(number).value;
591+
positive: atomic_func!(|num, vm| Self::number_float(num, vm).to_pyresult(vm)),
592+
absolute: atomic_func!(|num, vm| {
593+
let value = Self::number_downcast(num).value;
570594
value.abs().to_pyresult(vm)
571595
}),
572-
boolean: Some(|number, _vm| Ok(Self::number_downcast(number).value.is_zero())),
573-
int: Some(|number, vm| {
574-
let value = Self::number_downcast(number).value;
596+
boolean: atomic_func!(|num, _vm| Ok(Self::number_downcast(num).value.is_zero())),
597+
int: atomic_func!(|num, vm| {
598+
let value = Self::number_downcast(num).value;
575599
try_to_bigint(value, vm).map(|x| vm.ctx.new_int(x))
576600
}),
577-
float: Some(|number, vm| Ok(Self::number_float(number, vm))),
578-
floor_divide: Some(|number, other, vm| {
579-
Self::number_general_op(number, other, inner_floordiv, vm)
601+
float: atomic_func!(|num, vm| Ok(Self::number_float(num, vm))),
602+
floor_divide: atomic_func!(|num, other, vm| {
603+
Self::number_general_op(num, other, inner_floordiv, vm)
580604
}),
581-
true_divide: Some(|number, other, vm| {
582-
Self::number_general_op(number, other, inner_div, vm)
605+
true_divide: atomic_func!(|num, other, vm| {
606+
Self::number_general_op(num, other, inner_div, vm)
583607
}),
584608
..PyNumberMethods::NOT_IMPLEMENTED
585609
};

vm/src/builtins/int.rs

Lines changed: 48 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use crate::{
1515
TryFromBorrowedObject, VirtualMachine,
1616
};
1717
use bstr::ByteSlice;
18+
use crossbeam_utils::atomic::AtomicCell;
1819
use num_bigint::{BigInt, BigUint, Sign};
1920
use num_integer::Integer;
2021
use num_traits::{One, Pow, PrimInt, Signed, ToPrimitive, Zero};
@@ -755,36 +756,59 @@ impl Hashable for PyInt {
755756
}
756757
}
757758

759+
macro_rules! atomic_func {
760+
($x:expr) => {
761+
AtomicCell::new(Some($x))
762+
};
763+
}
764+
758765
impl AsNumber for PyInt {
759766
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);
767+
add: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a + b, vm)),
768+
subtract: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a - b, vm)),
769+
multiply: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a * b, vm)),
770+
remainder: atomic_func!(|num, other, vm| Self::number_general_op(
771+
num, other, inner_mod, vm
772+
)),
773+
divmod: atomic_func!(|num, other, vm| Self::number_general_op(
774+
num,
775+
other,
776+
inner_divmod,
777+
vm
778+
)),
779+
power: atomic_func!(|num, other, vm| Self::number_general_op(num, other, inner_pow, vm)),
780+
negative: atomic_func!(|num, vm| (&Self::number_downcast(num).value).neg().to_pyresult(vm)),
781+
positive: atomic_func!(|num, vm| Ok(Self::number_int(num, vm).into())),
782+
absolute: atomic_func!(|num, vm| Self::number_downcast(num).value.abs().to_pyresult(vm)),
783+
boolean: atomic_func!(|num, _vm| Ok(Self::number_downcast(num).value.is_zero())),
784+
invert: atomic_func!(|num, vm| (&Self::number_downcast(num).value).not().to_pyresult(vm)),
785+
lshift: atomic_func!(|num, other, vm| Self::number_general_op(
786+
num,
787+
other,
788+
inner_lshift,
789+
vm
790+
)),
791+
rshift: atomic_func!(|num, other, vm| Self::number_general_op(
792+
num,
793+
other,
794+
inner_rshift,
795+
vm
796+
)),
797+
and: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a & b, vm)),
798+
xor: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a ^ b, vm)),
799+
or: atomic_func!(|num, other, vm| Self::number_int_op(num, other, |a, b| a | b, vm)),
800+
int: atomic_func!(|num, other| Ok(Self::number_int(num, other))),
801+
float: atomic_func!(|num, vm| {
802+
let zelf = Self::number_downcast(num);
779803
try_to_float(&zelf.value, vm).map(|x| vm.ctx.new_float(x))
780804
}),
781-
floor_divide: Some(|number, other, vm| {
782-
Self::number_general_op(number, other, inner_floordiv, vm)
805+
floor_divide: atomic_func!(|num, other, vm| {
806+
Self::number_general_op(num, other, inner_floordiv, vm)
783807
}),
784-
true_divide: Some(|number, other, vm| {
785-
Self::number_general_op(number, other, inner_truediv, vm)
808+
true_divide: atomic_func!(|num, other, vm| {
809+
Self::number_general_op(num, other, inner_truediv, vm)
786810
}),
787-
index: Some(|number, vm| Ok(Self::number_int(number, vm))),
811+
index: atomic_func!(|num, vm| Ok(Self::number_int(num, vm))),
788812
..PyNumberMethods::NOT_IMPLEMENTED
789813
};
790814
}

0 commit comments

Comments
 (0)