From f1dc0a6fcbb66bd2029d2f14670c87e3e8bd6e68 Mon Sep 17 00:00:00 2001 From: Zhiyan Xiao Date: Mon, 6 Mar 2023 19:17:36 +0900 Subject: [PATCH] Improve: binary ops with Number Protocol --- vm/src/protocol/mod.rs | 2 +- vm/src/protocol/number.rs | 121 +++---- vm/src/stdlib/builtins.rs | 6 +- vm/src/vm/vm_ops.rs | 689 ++++++++++---------------------------- 4 files changed, 221 insertions(+), 597 deletions(-) diff --git a/vm/src/protocol/mod.rs b/vm/src/protocol/mod.rs index f355b61ea..7cde1a875 100644 --- a/vm/src/protocol/mod.rs +++ b/vm/src/protocol/mod.rs @@ -10,5 +10,5 @@ pub use buffer::{BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, V pub use callable::PyCallable; pub use iter::{PyIter, PyIterIter, PyIterReturn}; pub use mapping::{PyMapping, PyMappingMethods}; -pub use number::{PyNumber, PyNumberMethods, PyNumberMethodsOffset}; +pub use number::{PyNumber, PyNumberBinaryOpSlot, PyNumberMethods}; pub use sequence::{PySequence, PySequenceMethods}; diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs index 8abc54e1b..a7e9efd56 100644 --- a/vm/src/protocol/number.rs +++ b/vm/src/protocol/number.rs @@ -8,7 +8,6 @@ use crate::{ VirtualMachine, }; use crossbeam_utils::atomic::AtomicCell; -use std::ptr; type UnaryFunc = AtomicCell PyResult>>; type BinaryFunc = @@ -110,7 +109,6 @@ impl PyObject { } #[derive(Default)] -// #[repr(C)] pub struct PyNumberMethods { /* Number implementations must check *both* arguments for proper type and implement the necessary conversions @@ -138,7 +136,6 @@ pub struct PyNumberMethods { pub inplace_subtract: BinaryFunc, pub inplace_multiply: BinaryFunc, pub inplace_remainder: BinaryFunc, - pub inplace_divmod: BinaryFunc, pub inplace_power: BinaryFunc, pub inplace_lshift: BinaryFunc, pub inplace_rshift: BinaryFunc, @@ -184,7 +181,6 @@ impl PyNumberMethods { inplace_subtract: AtomicCell::new(None), inplace_multiply: AtomicCell::new(None), inplace_remainder: AtomicCell::new(None), - inplace_divmod: AtomicCell::new(None), inplace_power: AtomicCell::new(None), inplace_lshift: AtomicCell::new(None), inplace_rshift: AtomicCell::new(None), @@ -199,32 +195,58 @@ impl PyNumberMethods { matrix_multiply: AtomicCell::new(None), inplace_matrix_multiply: AtomicCell::new(None), }; + + pub fn get_binary_op(&self, op_slot: &PyNumberBinaryOpSlot) -> PyResult<&BinaryFunc> { + use PyNumberBinaryOpSlot::*; + let binary_op = match op_slot { + Add => &self.add, + Subtract => &self.subtract, + Multiply => &self.multiply, + Remainder => &self.remainder, + Divmod => &self.divmod, + Power => &self.power, + Lshift => &self.lshift, + Rshift => &self.rshift, + And => &self.and, + Xor => &self.xor, + Or => &self.or, + InplaceAdd => &self.inplace_add, + InplaceSubtract => &self.inplace_subtract, + InplaceMultiply => &self.inplace_multiply, + InplaceRemainder => &self.inplace_remainder, + InplacePower => &self.inplace_power, + InplaceLshift => &self.inplace_lshift, + InplaceRshift => &self.inplace_rshift, + InplaceAnd => &self.inplace_and, + InplaceXor => &self.inplace_xor, + InplaceOr => &self.inplace_or, + FloorDivide => &self.floor_divide, + TrueDivide => &self.true_divide, + InplaceFloorDivide => &self.inplace_floor_divide, + InplaceTrueDivide => &self.inplace_true_divide, + MatrixMultiply => &self.matrix_multiply, + InplaceMatrixMultiply => &self.inplace_matrix_multiply, + }; + Ok(binary_op) + } } -pub enum PyNumberMethodsOffset { +pub enum PyNumberBinaryOpSlot { Add, Subtract, Multiply, Remainder, Divmod, Power, - Negative, - Positive, - Absolute, - Boolean, - Invert, Lshift, Rshift, And, Xor, Or, - Int, - Float, InplaceAdd, InplaceSubtract, InplaceMultiply, InplaceRemainder, - InplaceDivmod, InplacePower, InplaceLshift, InplaceRshift, @@ -235,63 +257,10 @@ pub enum PyNumberMethodsOffset { TrueDivide, InplaceFloorDivide, InplaceTrueDivide, - Index, MatrixMultiply, InplaceMatrixMultiply, } -impl PyNumberMethodsOffset { - pub fn method(&self, methods: &PyNumberMethods, vm: &VirtualMachine) -> PyResult<&BinaryFunc> { - use PyNumberMethodsOffset::*; - unsafe { - match self { - // BinaryFunc - Add => ptr::addr_of!(methods.add), - Subtract => ptr::addr_of!(methods.subtract), - Multiply => ptr::addr_of!(methods.multiply), - Remainder => ptr::addr_of!(methods.remainder), - Divmod => ptr::addr_of!(methods.divmod), - Power => ptr::addr_of!(methods.power), - Lshift => ptr::addr_of!(methods.lshift), - Rshift => ptr::addr_of!(methods.rshift), - And => ptr::addr_of!(methods.and), - Xor => ptr::addr_of!(methods.xor), - Or => ptr::addr_of!(methods.or), - InplaceAdd => ptr::addr_of!(methods.inplace_add), - InplaceSubtract => ptr::addr_of!(methods.inplace_subtract), - InplaceMultiply => ptr::addr_of!(methods.inplace_multiply), - InplaceRemainder => ptr::addr_of!(methods.inplace_remainder), - InplaceDivmod => ptr::addr_of!(methods.inplace_divmod), - InplacePower => ptr::addr_of!(methods.inplace_power), - InplaceLshift => ptr::addr_of!(methods.inplace_lshift), - InplaceRshift => ptr::addr_of!(methods.inplace_rshift), - InplaceAnd => ptr::addr_of!(methods.inplace_and), - InplaceXor => ptr::addr_of!(methods.inplace_xor), - InplaceOr => ptr::addr_of!(methods.inplace_or), - FloorDivide => ptr::addr_of!(methods.floor_divide), - TrueDivide => ptr::addr_of!(methods.true_divide), - InplaceFloorDivide => ptr::addr_of!(methods.inplace_floor_divide), - InplaceTrueDivide => ptr::addr_of!(methods.inplace_true_divide), - MatrixMultiply => ptr::addr_of!(methods.matrix_multiply), - InplaceMatrixMultiply => ptr::addr_of!(methods.inplace_matrix_multiply), - // UnaryFunc - Negative => ptr::null(), - Positive => ptr::null(), - Absolute => ptr::null(), - Boolean => ptr::null(), - Invert => ptr::null(), - Int => ptr::null(), - Float => ptr::null(), - Index => ptr::null(), - } - .as_ref() - .ok_or_else(|| { - vm.new_value_error("No unaryop supported for PyNumberMethodsOffset".to_owned()) - }) - } - } -} - #[derive(Copy, Clone)] pub struct PyNumber<'a> { pub obj: &'a PyObject, @@ -314,12 +283,12 @@ impl PyNumber<'_> { obj.class().mro_find_map(|x| x.slots.as_number.load()) } - pub fn methods<'a>( - &'a self, - op_slot: &'a PyNumberMethodsOffset, - vm: &VirtualMachine, - ) -> PyResult<&BinaryFunc> { - op_slot.method(self.methods, vm) + pub fn methods(&self) -> &PyNumberMethods { + self.methods + } + + pub fn get_binary_op(&self, op_slot: &PyNumberBinaryOpSlot) -> PyResult<&BinaryFunc> { + self.methods().get_binary_op(op_slot) } // PyNumber_Check @@ -336,12 +305,12 @@ impl PyNumber<'_> { // PyIndex_Check pub fn is_index(&self) -> bool { - self.methods.index.load().is_some() + self.methods().index.load().is_some() } #[inline] pub fn int(self, vm: &VirtualMachine) -> Option> { - self.methods.int.load().map(|f| { + self.methods().int.load().map(|f| { let ret = f(self, vm)?; let value = if !ret.class().is(PyInt::class(vm)) { warnings::warn( @@ -365,7 +334,7 @@ impl PyNumber<'_> { #[inline] pub fn index(self, vm: &VirtualMachine) -> Option> { - self.methods.index.load().map(|f| { + self.methods().index.load().map(|f| { let ret = f(self, vm)?; let value = if !ret.class().is(PyInt::class(vm)) { warnings::warn( @@ -389,7 +358,7 @@ impl PyNumber<'_> { #[inline] pub fn float(self, vm: &VirtualMachine) -> Option>> { - self.methods.float.load().map(|f| { + self.methods().float.load().map(|f| { let ret = f(self, vm)?; let value = if !ret.class().is(PyFloat::class(vm)) { warnings::warn( diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs index 4ebdbd4f0..e382546b6 100644 --- a/vm/src/stdlib/builtins.rs +++ b/vm/src/stdlib/builtins.rs @@ -21,7 +21,7 @@ mod builtins { ArgBytesLike, ArgCallable, ArgIntoBool, ArgIterable, ArgMapping, ArgStrOrBytesLike, Either, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs, PyArithmeticValue, }, - protocol::{PyIter, PyIterReturn, PyNumberMethodsOffset}, + protocol::{PyIter, PyIterReturn, PyNumberBinaryOpSlot}, py_io, readline::{Readline, ReadlineResult}, stdlib::sys, @@ -605,9 +605,7 @@ mod builtins { modulus, } = args; match modulus { - None => vm.binary_op(&x, &y, PyNumberMethodsOffset::Power, "pow", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }), + None => vm.binary_op(&x, &y, &PyNumberBinaryOpSlot::Power, "pow"), Some(z) => { let try_pow_value = |obj: &PyObject, args: (PyObjectRef, PyObjectRef, PyObjectRef)| diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs index c7a281803..fb524bc16 100644 --- a/vm/src/vm/vm_ops.rs +++ b/vm/src/vm/vm_ops.rs @@ -1,13 +1,34 @@ use super::{PyMethod, VirtualMachine}; use crate::{ - builtins::{PyInt, PyIntRef, PyStr, PyStrInterned, PyStrRef}, - function::PyArithmeticValue, + builtins::{PyInt, PyIntRef, PyStr, PyStrRef}, object::{AsObject, PyObject, PyObjectRef, PyResult}, - protocol::{PyIterReturn, PyNumberMethodsOffset, PySequence}, + protocol::{PyIterReturn, PyNumberBinaryOpSlot, PySequence}, types::PyComparisonOp, }; use num_traits::ToPrimitive; +macro_rules! binary_func { + ($fn:ident, $op_slot:ident, $op:expr) => { + pub fn $fn(&self, a: &PyObject, b: &PyObject) -> PyResult { + self.binary_op(a, b, &PyNumberBinaryOpSlot::$op_slot, $op) + } + }; +} + +macro_rules! inplace_binary_func { + ($fn:ident, $iop_slot:ident, $op_slot:ident, $op:expr) => { + pub fn $fn(&self, a: &PyObject, b: &PyObject) -> PyResult { + self.binary_iop( + a, + b, + &PyNumberBinaryOpSlot::$iop_slot, + &PyNumberBinaryOpSlot::$op_slot, + $op, + ) + } + }; +} + /// Collection of operators impl VirtualMachine { #[inline] @@ -104,585 +125,221 @@ impl VirtualMachine { } } - // TODO: Should be deleted after transplanting complete number protocol - /// Calls a method on `obj` passing `arg`, if the method exists. + /// Calling scheme used for binary operations: /// - /// Otherwise, or if the result is the special `NotImplemented` built-in constant, - /// calls `unsupported` to determine fallback value. - pub fn call_or_unsupported( - &self, - obj: &PyObject, - arg: &PyObject, - method: &'static PyStrInterned, - unsupported: F, - ) -> PyResult - where - F: Fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - { - if let Some(method_or_err) = self.get_method(obj.to_owned(), method) { - let method = method_or_err?; - let result = method.call((arg.to_owned(),), self)?; - if let PyArithmeticValue::Implemented(x) = PyArithmeticValue::from_object(self, result) - { - return Ok(x); - } - } - unsupported(self, obj, arg) - } - - // TODO: Should be deleted after transplanting complete number protocol - /// Calls a method, falling back to its reflection with the operands - /// reversed, and then to the value provided by `unsupported`. + /// Order operations are tried until either a valid result or error: + /// b.op(a,b)[*], a.op(a,b), b.op(a,b) /// - /// For example: the following: - /// - /// `call_or_reflection(lhs, rhs, "__and__", "__rand__", unsupported)` - /// - /// 1. Calls `__and__` with `lhs` and `rhs`. - /// 2. If above is not implemented, calls `__rand__` with `rhs` and `lhs`. - /// 3. If above is not implemented, invokes `unsupported` for the result. - pub fn call_or_reflection( - &self, - lhs: &PyObject, - rhs: &PyObject, - default: &'static PyStrInterned, - reflection: &'static PyStrInterned, - unsupported: fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - ) -> PyResult { - if rhs.fast_isinstance(lhs.class()) { - let lop = lhs.get_class_attr(reflection); - let rop = rhs.get_class_attr(reflection); - if let Some((lop, rop)) = lop.zip(rop) { - if !lop.is(&rop) { - if let Ok(r) = self.call_or_unsupported(rhs, lhs, reflection, |vm, _, _| { - Err(vm.new_exception_empty(vm.ctx.exceptions.exception_type.to_owned())) - }) { - return Ok(r); - } - } - } - } - // Try to call the default method - self.call_or_unsupported(lhs, rhs, default, move |vm, lhs, rhs| { - // Try to call the reflection method - // don't call reflection method if operands are of the same type - if !lhs.class().is(rhs.class()) { - vm.call_or_unsupported(rhs, lhs, reflection, |_, rhs, lhs| { - // switch them around again - unsupported(vm, lhs, rhs) - }) - } else { - unsupported(vm, lhs, rhs) - } - }) - } - - fn binary_op1( - &self, - a: &PyObject, - b: &PyObject, - op_slot: PyNumberMethodsOffset, - unsupported: F, - ) -> PyResult - where - F: Fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - { + /// [*] only when Py_TYPE(a) != Py_TYPE(b) && Py_TYPE(b) is a subclass of Py_TYPE(a) + fn binary_op1(&self, a: &PyObject, b: &PyObject, op_slot: &PyNumberBinaryOpSlot) -> PyResult { let num_a = a.to_number(); let num_b = b.to_number(); - let slot_a = num_a.methods(&op_slot, self)?.load(); - let slot_b = num_b.methods(&op_slot, self)?.load(); + let slot_a = num_a.get_binary_op(op_slot)?.load(); + let mut slot_b = if b.class().is(a.class()) { + None + } else { + match num_b.get_binary_op(op_slot)?.load() { + Some(slot_b) + if slot_b as usize == slot_a.map(|s| s as usize).unwrap_or_default() => + { + None + } + slot_b => slot_b, + } + }; if let Some(slot_a) = slot_a { - if let Some(slot_b) = slot_b { - // Check if `a` is subclass of `b` + if let Some(slot_bb) = slot_b { if b.fast_isinstance(a.class()) { - let ret = slot_b(num_a, b, self)?; - if ret.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Ne, - self, - )? { - return Ok(ret); + let x = slot_bb(num_a, b, self)?; + if !x.is(&self.ctx.not_implemented) { + return Ok(x); } + slot_b = None; } } - - let ret = slot_a(num_a, b, self)?; - if ret.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Ne, - self, - )? { - return Ok(ret); + let x = slot_a(num_a, b, self)?; + if !x.is(&self.ctx.not_implemented) { + return Ok(x); } } - // No slot_a or Not implemented if let Some(slot_b) = slot_b { - let ret = slot_b(num_a, b, self)?; - if ret.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Ne, - self, - )? { - return Ok(ret); + let x = slot_b(num_a, b, self)?; + if !x.is(&self.ctx.not_implemented) { + return Ok(x); } } - // Both slot_a & slot_b don't exist or are not implemented. - unsupported(self, a, b) + Ok(self.ctx.not_implemented()) } - /// `binary_op()` can work only with [`PyNumberMethods::BinaryFunc`]. - pub fn binary_op( + pub fn binary_op( &self, a: &PyObject, b: &PyObject, - op_slot: PyNumberMethodsOffset, + op_slot: &PyNumberBinaryOpSlot, op: &str, - unsupported: F, - ) -> PyResult - where - F: Fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - { - let result = self.binary_op1(a, b, op_slot, unsupported)?; - - if result.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Eq, - self, - )? { - Err(self.new_unsupported_binop_error(a, b, op)) - } else { - Ok(result) + ) -> PyResult { + let result = self.binary_op1(a, b, op_slot)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); } + Err(self.new_unsupported_binop_error(a, b, op)) } - /// ### Binary in-place operators + /// Binary in-place operators /// /// The in-place operators are defined to fall back to the 'normal', /// non in-place operations, if the in-place methods are not in place. /// /// - If the left hand object has the appropriate struct members, and - /// they are filled, call the appropriate function and return the - /// result. No coercion is done on the arguments; the left-hand object - /// is the one the operation is performed on, and it's up to the - /// function to deal with the right-hand object. + /// they are filled, call the appropriate function and return the + /// result. No coercion is done on the arguments; the left-hand object + /// is the one the operation is performed on, and it's up to the + /// function to deal with the right-hand object. /// /// - Otherwise, in-place modification is not supported. Handle it exactly as - /// a non in-place operation of the same kind. - fn binary_iop1( + /// a non in-place operation of the same kind. + fn binary_iop1( &self, a: &PyObject, b: &PyObject, - iop_slot: PyNumberMethodsOffset, - op_slot: PyNumberMethodsOffset, - unsupported: F, - ) -> PyResult - where - F: Fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - { + iop_slot: &PyNumberBinaryOpSlot, + op_slot: &PyNumberBinaryOpSlot, + ) -> PyResult { let num_a = a.to_number(); - let slot_a = num_a.methods(&iop_slot, self)?.load(); - - if let Some(slot_a) = slot_a { - let ret = slot_a(num_a, b, self)?; - if ret.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Ne, - self, - )? { - return Ok(ret); + if let Some(slot) = num_a.get_binary_op(iop_slot)?.load() { + let x = slot(num_a, b, self)?; + if !x.is(&self.ctx.not_implemented) { + return Ok(x); } } - - self.binary_op1(a, b, op_slot, unsupported) + self.binary_op1(a, b, op_slot) } - /// `binary_iop()` can work only with [`PyNumberMethods::BinaryFunc`]. - fn binary_iop( + fn binary_iop( &self, a: &PyObject, b: &PyObject, - iop_slot: PyNumberMethodsOffset, - op_slot: PyNumberMethodsOffset, + iop_slot: &PyNumberBinaryOpSlot, + op_slot: &PyNumberBinaryOpSlot, op: &str, - unsupported: F, - ) -> PyResult - where - F: Fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult, - { - let result = self.binary_iop1(a, b, iop_slot, op_slot, unsupported)?; - - if result.rich_compare_bool( - self.ctx.not_implemented.as_object(), - PyComparisonOp::Eq, - self, - )? { - Err(self.new_unsupported_binop_error(a, b, op)) - } else { - Ok(result) + ) -> PyResult { + let result = self.binary_iop1(a, b, iop_slot, op_slot)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); } + Err(self.new_unsupported_binop_error(a, b, op)) } + binary_func!(_sub, Subtract, "-"); + binary_func!(_mod, Remainder, "%"); + binary_func!(_divmod, Divmod, "divmod"); + binary_func!(_pow, Power, "**"); + binary_func!(_lshift, Lshift, "<<"); + binary_func!(_rshift, Rshift, ">>"); + binary_func!(_and, And, "&"); + binary_func!(_xor, Xor, "^"); + binary_func!(_or, Or, "|"); + binary_func!(_floordiv, FloorDivide, "//"); + binary_func!(_truediv, TrueDivide, "/"); + binary_func!(_matmul, MatrixMultiply, "@"); + + inplace_binary_func!(_isub, InplaceSubtract, Subtract, "-="); + inplace_binary_func!(_imod, InplaceRemainder, Remainder, "%="); + inplace_binary_func!(_ipow, InplacePower, Power, "**="); + inplace_binary_func!(_ilshift, InplaceLshift, Lshift, "<<="); + inplace_binary_func!(_irshift, InplaceRshift, Rshift, ">>="); + inplace_binary_func!(_iand, InplaceAnd, And, "&="); + inplace_binary_func!(_ixor, InplaceXor, Xor, "^="); + inplace_binary_func!(_ior, InplaceOr, Or, "|="); + inplace_binary_func!(_ifloordiv, InplaceFloorDivide, FloorDivide, "//="); + inplace_binary_func!(_itruediv, InplaceTrueDivide, TrueDivide, "/="); + inplace_binary_func!(_imatmul, InplaceMatrixMultiply, MatrixMultiply, "@="); + pub fn _add(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Add, "+", |vm, a, b| { - let seq_a = PySequence::try_protocol(a, vm); - - if let Ok(seq_a) = seq_a { - let ret = seq_a.concat(b, vm)?; - if ret.rich_compare_bool( - vm.ctx.not_implemented.as_object(), - PyComparisonOp::Ne, - vm, - )? { - return Ok(ret); - } + let result = self.binary_op1(a, b, &PyNumberBinaryOpSlot::Add)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); + } + if let Ok(seq_a) = PySequence::try_protocol(a, self) { + let result = seq_a.concat(b, self)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); } - - Ok(vm.ctx.not_implemented()) - }) + } + Err(self.new_unsupported_binop_error(a, b, "+")) } pub fn _iadd(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( + let result = self.binary_iop1( a, b, - PyNumberMethodsOffset::InplaceAdd, - PyNumberMethodsOffset::Add, - "+=", - |vm, a, b| { - let seq_a = PySequence::try_protocol(a, vm); - - if let Ok(seq_a) = seq_a { - return seq_a.inplace_concat(b, vm); - } - - Ok(vm.ctx.not_implemented()) - }, - ) - } - - pub fn _sub(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Subtract, "-", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _isub(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceSubtract, - PyNumberMethodsOffset::Subtract, - "-=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) + &PyNumberBinaryOpSlot::InplaceAdd, + &PyNumberBinaryOpSlot::Add, + )?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); + } + if let Ok(seq_a) = PySequence::try_protocol(a, self) { + let result = seq_a.inplace_concat(b, self)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); + } + } + Err(self.new_unsupported_binop_error(a, b, "+=")) } pub fn _mul(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Multiply, "*", |vm, a, b| { - // TODO: check if PySequence::with_methods can replace try_protocol - let seq_a = PySequence::try_protocol(a, vm); - let seq_b = PySequence::try_protocol(b, vm); - - // TODO: I think converting to isize process should be handled in repeat function. - // TODO: This can be helpful to unify the sequence protocol's closure. - - if let Ok(seq_a) = seq_a { - let n = b.try_int(vm)?.as_bigint().to_isize().ok_or_else(|| { - vm.new_overflow_error("repeated bytes are too long".to_owned()) + let result = self.binary_op1(a, b, &PyNumberBinaryOpSlot::Multiply)?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); + } + if let Ok(seq_a) = PySequence::try_protocol(a, self) { + let n = + b.try_index(self)?.as_bigint().to_isize().ok_or_else(|| { + self.new_overflow_error("repeated bytes are too long".to_owned()) })?; - - return seq_a.repeat(n, vm); - } else if let Ok(seq_b) = seq_b { - let n = a.try_int(vm)?.as_bigint().to_isize().ok_or_else(|| { - vm.new_overflow_error("repeated bytes are too long".to_owned()) + return seq_a.repeat(n, self); + } else if let Ok(seq_b) = PySequence::try_protocol(b, self) { + let n = + a.try_index(self)?.as_bigint().to_isize().ok_or_else(|| { + self.new_overflow_error("repeated bytes are too long".to_owned()) })?; - - return seq_b.repeat(n, vm); - } - - Ok(vm.ctx.not_implemented()) - }) + return seq_b.repeat(n, self); + } + Err(self.new_unsupported_binop_error(a, b, "*")) } pub fn _imul(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( + let result = self.binary_iop1( a, b, - PyNumberMethodsOffset::InplaceMultiply, - PyNumberMethodsOffset::Multiply, - "*=", - |vm, a, b| { - // TODO: check if PySequence::with_methods can replace try_protocol - let seq_a = PySequence::try_protocol(a, vm); - let seq_b = PySequence::try_protocol(b, vm); - - if let Ok(seq_a) = seq_a { - let n = b.try_int(vm)?.as_bigint().to_isize().ok_or_else(|| { - vm.new_overflow_error("repeated bytes are too long".to_owned()) - })?; - - return seq_a.inplace_repeat(n, vm); - } else if let Ok(seq_b) = seq_b { - let n = a.try_int(vm)?.as_bigint().to_isize().ok_or_else(|| { - vm.new_overflow_error("repeated bytes are too long".to_owned()) - })?; - - /* Note that the right hand operand should not be - * mutated in this case so sq_inplace_repeat is not - * used. */ - return seq_b.repeat(n, vm); - } - - Ok(vm.ctx.not_implemented()) - }, - ) - } - - pub fn _matmul(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op( - a, - b, - PyNumberMethodsOffset::MatrixMultiply, - "@", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _imatmul(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceMatrixMultiply, - PyNumberMethodsOffset::MatrixMultiply, - "@=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _truediv(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::TrueDivide, "/", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _itruediv(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceTrueDivide, - PyNumberMethodsOffset::TrueDivide, - "/=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _floordiv(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op( - a, - b, - PyNumberMethodsOffset::FloorDivide, - "//", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _ifloordiv(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceFloorDivide, - PyNumberMethodsOffset::FloorDivide, - "//=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _pow(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Power, "**", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _ipow(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplacePower, - PyNumberMethodsOffset::Power, - "**=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - // TODO: `str` modular opertation(mod, imod) is not supported now. Should implement it. - pub fn _mod(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_reflection( - a, - b, - identifier!(self, __mod__), - identifier!(self, __rmod__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "%")), - ) - - // self.binary_op(a, b, PyNumberMethodsOffset::Remainder, "%", |vm, _, _| { - // Ok(vm.ctx.not_implemented()) - // }) - } - - pub fn _imod(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_unsupported(a, b, identifier!(self, __imod__), |vm, a, b| { - vm.call_or_reflection( - a, - b, - identifier!(self, __mod__), - identifier!(self, __rmod__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "%=")), - ) - }) - - // self.binary_iop( - // a, - // b, - // PyNumberMethodsOffset::InplaceRemainder, - // PyNumberMethodsOffset::Remainder, - // "%=", - // |vm, _, _| Ok(vm.ctx.not_implemented()), - // ) - } - - pub fn _divmod(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Divmod, "divmod", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _lshift(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Lshift, "<<", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _ilshift(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceLshift, - PyNumberMethodsOffset::Lshift, - "<<=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _rshift(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Rshift, ">>", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _irshift(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceRshift, - PyNumberMethodsOffset::Rshift, - ">>=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - pub fn _xor(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_op(a, b, PyNumberMethodsOffset::Xor, "^", |vm, _, _| { - Ok(vm.ctx.not_implemented()) - }) - } - - pub fn _ixor(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.binary_iop( - a, - b, - PyNumberMethodsOffset::InplaceXor, - PyNumberMethodsOffset::Xor, - "^=", - |vm, _, _| Ok(vm.ctx.not_implemented()), - ) - } - - // TODO: `or` method doesn't work because of structure of `type_::or_()`. - // TODO: It should be changed by adjusting with AsNumber. - pub fn _or(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_reflection( - a, - b, - identifier!(self, __or__), - identifier!(self, __ror__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "|")), - ) - - // self.binary_op(a, b, PyNumberMethodsOffset::Or, "|", |vm, _, _| { - // Ok(vm.ctx.not_implemented()) - // }) - } - - pub fn _ior(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_unsupported(a, b, identifier!(self, __ior__), |vm, a, b| { - vm.call_or_reflection( - a, - b, - identifier!(self, __or__), - identifier!(self, __ror__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "|=")), - ) - }) - - // self.binary_iop( - // a, - // b, - // PyNumberMethodsOffset::InplaceOr, - // PyNumberMethodsOffset::Or, - // "|=", - // |vm, _, _| Ok(vm.ctx.not_implemented()), - // ) - } - - // TODO: `and` method doesn't work because of structure of `set`. - // TODO: It should be changed by adjusting with AsNumber. - pub fn _and(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_reflection( - a, - b, - identifier!(self, __and__), - identifier!(self, __rand__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "&")), - ) - - // self.binary_op(a, b, PyNumberMethodsOffset::And, "&", |vm, _, _| { - // Ok(vm.ctx.not_implemented()) - // }) - } - - pub fn _iand(&self, a: &PyObject, b: &PyObject) -> PyResult { - self.call_or_unsupported(a, b, identifier!(self, __iand__), |vm, a, b| { - vm.call_or_reflection( - a, - b, - identifier!(self, __and__), - identifier!(self, __rand__), - |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "&=")), - ) - }) - - // self.binary_iop( - // a, - // b, - // PyNumberMethodsOffset::InplaceAnd, - // PyNumberMethodsOffset::And, - // "&=", - // |vm, _, _| Ok(vm.ctx.not_implemented()), - // ) + &PyNumberBinaryOpSlot::InplaceMultiply, + &PyNumberBinaryOpSlot::Multiply, + )?; + if !result.is(&self.ctx.not_implemented) { + return Ok(result); + } + if let Ok(seq_a) = PySequence::try_protocol(a, self) { + let n = + b.try_index(self)?.as_bigint().to_isize().ok_or_else(|| { + self.new_overflow_error("repeated bytes are too long".to_owned()) + })?; + return seq_a.inplace_repeat(n, self); + } else if let Ok(seq_b) = PySequence::try_protocol(b, self) { + let n = + a.try_index(self)?.as_bigint().to_isize().ok_or_else(|| { + self.new_overflow_error("repeated bytes are too long".to_owned()) + })?; + /* Note that the right hand operand should not be + * mutated in this case so inplace_repeat is not + * used. */ + return seq_b.repeat(n, self); + } + Err(self.new_unsupported_binop_error(a, b, "*=")) } pub fn _abs(&self, a: &PyObject) -> PyResult {