Skip to content

Commit eb41b6a

Browse files
Karatussxiaozhiyan
Karatuss
authored and
xiaozhiyan
committed
Binary ops with number protocol (partially imported)
from RustPython#4139 RustPython#4139
1 parent 8d6d47a commit eb41b6a

4 files changed

Lines changed: 479 additions & 170 deletions

File tree

‎vm/src/protocol/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,5 +8,5 @@ mod sequence;
88
pub use buffer::{BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, VecBuffer};
99
pub use iter::{PyIter, PyIterIter, PyIterReturn};
1010
pub use mapping::{PyMapping, PyMappingMethods};
11-
pub use number::{PyNumber, PyNumberMethods};
11+
pub use number::{PyNumber, PyNumberMethods, PyNumberMethodsOffset};
1212
pub use sequence::{PySequence, PySequenceMethods};

‎vm/src/protocol/number.rs‎

Lines changed: 104 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::{
88
VirtualMachine,
99
};
1010
use crossbeam_utils::atomic::AtomicCell;
11+
use std::ptr;
1112

1213
type UnaryFunc<R = PyObjectRef> = AtomicCell<Option<fn(PyNumber, &VirtualMachine) -> PyResult<R>>>;
1314
type BinaryFunc<R = PyObjectRef> =
@@ -109,6 +110,7 @@ impl PyObject {
109110
}
110111

111112
#[derive(Default)]
113+
// #[repr(C)]
112114
pub struct PyNumberMethods {
113115
/* Number implementations must check *both*
114116
arguments for proper type and implement the necessary conversions
@@ -199,6 +201,98 @@ impl PyNumberMethods {
199201
};
200202
}
201203

204+
pub enum PyNumberMethodsOffset {
205+
Add,
206+
Subtract,
207+
Multiply,
208+
Remainder,
209+
Divmod,
210+
Power,
211+
Negative,
212+
Positive,
213+
Absolute,
214+
Boolean,
215+
Invert,
216+
Lshift,
217+
Rshift,
218+
And,
219+
Xor,
220+
Or,
221+
Int,
222+
Float,
223+
InplaceAdd,
224+
InplaceSubtract,
225+
InplaceMultiply,
226+
InplaceRemainder,
227+
InplaceDivmod,
228+
InplacePower,
229+
InplaceLshift,
230+
InplaceRshift,
231+
InplaceAnd,
232+
InplaceXor,
233+
InplaceOr,
234+
FloorDivide,
235+
TrueDivide,
236+
InplaceFloorDivide,
237+
InplaceTrueDivide,
238+
Index,
239+
MatrixMultiply,
240+
InplaceMatrixMultiply,
241+
}
242+
243+
impl PyNumberMethodsOffset {
244+
pub fn method(&self, methods: &PyNumberMethods, vm: &VirtualMachine) -> PyResult<&BinaryFunc> {
245+
use PyNumberMethodsOffset::*;
246+
unsafe {
247+
match self {
248+
// BinaryFunc
249+
Add => ptr::addr_of!(methods.add),
250+
Subtract => ptr::addr_of!(methods.subtract),
251+
Multiply => ptr::addr_of!(methods.multiply),
252+
Remainder => ptr::addr_of!(methods.remainder),
253+
Divmod => ptr::addr_of!(methods.divmod),
254+
Power => ptr::addr_of!(methods.power),
255+
Lshift => ptr::addr_of!(methods.lshift),
256+
Rshift => ptr::addr_of!(methods.rshift),
257+
And => ptr::addr_of!(methods.and),
258+
Xor => ptr::addr_of!(methods.xor),
259+
Or => ptr::addr_of!(methods.or),
260+
InplaceAdd => ptr::addr_of!(methods.inplace_add),
261+
InplaceSubtract => ptr::addr_of!(methods.inplace_subtract),
262+
InplaceMultiply => ptr::addr_of!(methods.inplace_multiply),
263+
InplaceRemainder => ptr::addr_of!(methods.inplace_remainder),
264+
InplaceDivmod => ptr::addr_of!(methods.inplace_divmod),
265+
InplacePower => ptr::addr_of!(methods.inplace_power),
266+
InplaceLshift => ptr::addr_of!(methods.inplace_lshift),
267+
InplaceRshift => ptr::addr_of!(methods.inplace_rshift),
268+
InplaceAnd => ptr::addr_of!(methods.inplace_and),
269+
InplaceXor => ptr::addr_of!(methods.inplace_xor),
270+
InplaceOr => ptr::addr_of!(methods.inplace_or),
271+
FloorDivide => ptr::addr_of!(methods.floor_divide),
272+
TrueDivide => ptr::addr_of!(methods.true_divide),
273+
InplaceFloorDivide => ptr::addr_of!(methods.inplace_floor_divide),
274+
InplaceTrueDivide => ptr::addr_of!(methods.inplace_true_divide),
275+
MatrixMultiply => ptr::addr_of!(methods.matrix_multiply),
276+
InplaceMatrixMultiply => ptr::addr_of!(methods.inplace_matrix_multiply),
277+
// UnaryFunc
278+
Negative => ptr::null(),
279+
Positive => ptr::null(),
280+
Absolute => ptr::null(),
281+
Boolean => ptr::null(),
282+
Invert => ptr::null(),
283+
Int => ptr::null(),
284+
Float => ptr::null(),
285+
Index => ptr::null(),
286+
}
287+
.as_ref()
288+
.ok_or_else(|| {
289+
vm.new_value_error("No unaryop supported for PyNumberMethodsOffset".to_owned())
290+
})
291+
}
292+
}
293+
}
294+
295+
#[derive(Copy, Clone)]
202296
pub struct PyNumber<'a> {
203297
pub obj: &'a PyObject,
204298
methods: &'a PyNumberMethods,
@@ -220,8 +314,12 @@ impl PyNumber<'_> {
220314
obj.class().mro_find_map(|x| x.slots.as_number.load())
221315
}
222316

223-
pub fn methods(&self) -> &PyNumberMethods {
224-
self.methods
317+
pub fn methods<'a>(
318+
&'a self,
319+
op_slot: &'a PyNumberMethodsOffset,
320+
vm: &VirtualMachine,
321+
) -> PyResult<&BinaryFunc> {
322+
op_slot.method(self.methods, vm)
225323
}
226324

227325
// PyNumber_Check
@@ -238,12 +336,12 @@ impl PyNumber<'_> {
238336

239337
// PyIndex_Check
240338
pub fn is_index(&self) -> bool {
241-
self.methods().index.load().is_some()
339+
self.methods.index.load().is_some()
242340
}
243341

244342
#[inline]
245343
pub fn int(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
246-
self.methods().int.load().map(|f| {
344+
self.methods.int.load().map(|f| {
247345
let ret = f(self, vm)?;
248346
let value = if !ret.class().is(PyInt::class(vm)) {
249347
warnings::warn(
@@ -267,7 +365,7 @@ impl PyNumber<'_> {
267365

268366
#[inline]
269367
pub fn index(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
270-
self.methods().index.load().map(|f| {
368+
self.methods.index.load().map(|f| {
271369
let ret = f(self, vm)?;
272370
let value = if !ret.class().is(PyInt::class(vm)) {
273371
warnings::warn(
@@ -291,7 +389,7 @@ impl PyNumber<'_> {
291389

292390
#[inline]
293391
pub fn float(self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
294-
self.methods().float.load().map(|f| {
392+
self.methods.float.load().map(|f| {
295393
let ret = f(self, vm)?;
296394
let value = if !ret.class().is(PyFloat::class(vm)) {
297395
warnings::warn(

‎vm/src/stdlib/builtins.rs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ mod builtins {
2121
ArgBytesLike, ArgCallable, ArgIntoBool, ArgIterable, ArgMapping, ArgStrOrBytesLike,
2222
Either, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs, PyArithmeticValue,
2323
},
24-
protocol::{PyIter, PyIterReturn},
24+
protocol::{PyIter, PyIterReturn, PyNumberMethodsOffset},
2525
py_io,
2626
readline::{Readline, ReadlineResult},
2727
stdlib::sys,
@@ -605,13 +605,9 @@ mod builtins {
605605
modulus,
606606
} = args;
607607
match modulus {
608-
None => vm.call_or_reflection(
609-
&x,
610-
&y,
611-
identifier!(vm, __pow__),
612-
identifier!(vm, __rpow__),
613-
|vm, x, y| Err(vm.new_unsupported_binop_error(x, y, "pow")),
614-
),
608+
None => vm.binary_op(&x, &y, PyNumberMethodsOffset::Power, "pow", |vm, _, _| {
609+
Ok(vm.ctx.not_implemented())
610+
}),
615611
Some(z) => {
616612
let try_pow_value = |obj: &PyObject,
617613
args: (PyObjectRef, PyObjectRef, PyObjectRef)|

0 commit comments

Comments
 (0)