From 05900fbabfef4c7d09bc6d30f71e5a5d419a1fc3 Mon Sep 17 00:00:00 2001 From: snowapril Date: Tue, 28 Mar 2023 23:01:44 +0900 Subject: [PATCH] Fix power, inplace_power as ternary function Signed-off-by: snowapril --- vm/src/protocol/mod.rs | 2 +- vm/src/protocol/number.rs | 38 +++++++++--- vm/src/types/slot.rs | 20 +++++- vm/src/vm/vm_ops.rs | 124 +++++++++++++++++++++++++++++++++++++- 4 files changed, 167 insertions(+), 17 deletions(-) diff --git a/vm/src/protocol/mod.rs b/vm/src/protocol/mod.rs index dbbb39fbc..989cca73d 100644 --- a/vm/src/protocol/mod.rs +++ b/vm/src/protocol/mod.rs @@ -12,6 +12,6 @@ pub use iter::{PyIter, PyIterIter, PyIterReturn}; pub use mapping::{PyMapping, PyMappingMethods}; pub use number::{ PyNumber, PyNumberBinaryFunc, PyNumberBinaryOp, PyNumberMethods, PyNumberSlots, - PyNumberUnaryFunc, + PyNumberTernaryOp, PyNumberUnaryFunc, }; pub use sequence::{PySequence, PySequenceMethods}; diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs index f011376af..f49c1a54a 100644 --- a/vm/src/protocol/number.rs +++ b/vm/src/protocol/number.rs @@ -13,6 +13,7 @@ use crate::{ pub type PyNumberUnaryFunc = fn(&PyObject, &VirtualMachine) -> PyResult; pub type PyNumberBinaryFunc = fn(&PyObject, &PyObject, &VirtualMachine) -> PyResult; +pub type PyNumberTernaryFunc = fn(&PyObject, &PyObject, &PyObject, &VirtualMachine) -> PyResult; impl PyObject { #[inline] @@ -121,7 +122,7 @@ pub struct PyNumberMethods { pub multiply: Option, pub remainder: Option, pub divmod: Option, - pub power: Option, + pub power: Option, pub negative: Option, pub positive: Option, pub absolute: Option, @@ -139,7 +140,7 @@ pub struct PyNumberMethods { pub inplace_subtract: Option, pub inplace_multiply: Option, pub inplace_remainder: Option, - pub inplace_power: Option, + pub inplace_power: Option, pub inplace_lshift: Option, pub inplace_rshift: Option, pub inplace_and: Option, @@ -210,7 +211,6 @@ pub enum PyNumberBinaryOp { Multiply, Remainder, Divmod, - Power, Lshift, Rshift, And, @@ -220,7 +220,6 @@ pub enum PyNumberBinaryOp { InplaceSubtract, InplaceMultiply, InplaceRemainder, - InplacePower, InplaceLshift, InplaceRshift, InplaceAnd, @@ -234,6 +233,12 @@ pub enum PyNumberBinaryOp { InplaceMatrixMultiply, } +#[derive(Copy, Clone)] +pub enum PyNumberTernaryOp { + Power, + InplacePower, +} + #[derive(Default)] pub struct PyNumberSlots { pub add: AtomicCell>, @@ -241,7 +246,7 @@ pub struct PyNumberSlots { pub multiply: AtomicCell>, pub remainder: AtomicCell>, pub divmod: AtomicCell>, - pub power: AtomicCell>, + pub power: AtomicCell>, pub negative: AtomicCell>, pub positive: AtomicCell>, pub absolute: AtomicCell>, @@ -260,7 +265,7 @@ pub struct PyNumberSlots { pub right_multiply: AtomicCell>, pub right_remainder: AtomicCell>, pub right_divmod: AtomicCell>, - pub right_power: AtomicCell>, + pub right_power: AtomicCell>, pub right_lshift: AtomicCell>, pub right_rshift: AtomicCell>, pub right_and: AtomicCell>, @@ -271,7 +276,7 @@ pub struct PyNumberSlots { pub inplace_subtract: AtomicCell>, pub inplace_multiply: AtomicCell>, pub inplace_remainder: AtomicCell>, - pub inplace_power: AtomicCell>, + pub inplace_power: AtomicCell>, pub inplace_lshift: AtomicCell>, pub inplace_rshift: AtomicCell>, pub inplace_and: AtomicCell>, @@ -359,7 +364,6 @@ impl PyNumberSlots { Multiply => self.multiply.load(), Remainder => self.remainder.load(), Divmod => self.divmod.load(), - Power => self.power.load(), Lshift => self.lshift.load(), Rshift => self.rshift.load(), And => self.and.load(), @@ -369,7 +373,6 @@ impl PyNumberSlots { InplaceSubtract => self.inplace_subtract.load(), InplaceMultiply => self.inplace_multiply.load(), InplaceRemainder => self.inplace_remainder.load(), - InplacePower => self.inplace_power.load(), InplaceLshift => self.inplace_lshift.load(), InplaceRshift => self.inplace_rshift.load(), InplaceAnd => self.inplace_and.load(), @@ -392,7 +395,6 @@ impl PyNumberSlots { Multiply => self.right_multiply.load(), Remainder => self.right_remainder.load(), Divmod => self.right_divmod.load(), - Power => self.right_power.load(), Lshift => self.right_lshift.load(), Rshift => self.right_rshift.load(), And => self.right_and.load(), @@ -404,6 +406,22 @@ impl PyNumberSlots { _ => None, } } + + pub fn left_ternary_op(&self, op_slot: PyNumberTernaryOp) -> Option { + use PyNumberTernaryOp::*; + match op_slot { + Power => self.power.load(), + InplacePower => self.inplace_power.load(), + } + } + + pub fn right_ternary_op(&self, op_slot: PyNumberTernaryOp) -> Option { + use PyNumberTernaryOp::*; + match op_slot { + Power => self.right_power.load(), + _ => None, + } + } } #[derive(Copy, Clone)] pub struct PyNumber<'a>(&'a PyObject); diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs index 7b9cb64b6..5ccad5171 100644 --- a/vm/src/types/slot.rs +++ b/vm/src/types/slot.rs @@ -212,6 +212,20 @@ macro_rules! number_binary_right_op_wrapper { |a, b, vm| vm.call_special_method(b, identifier!(vm, $name), (a.to_owned(),)) }; } +macro_rules! number_ternary_op_wrapper { + ($name:ident) => { + |a, b, c, vm| { + vm.call_special_method(a, identifier!(vm, $name), (b.to_owned(), c.to_owned())) + } + }; +} +macro_rules! number_ternary_right_op_wrapper { + ($name:ident) => { + |a, b, c, vm| { + vm.call_special_method(b, identifier!(vm, $name), (a.to_owned(), c.to_owned())) + } + }; +} fn getitem_wrapper(obj: &PyObject, needle: K, vm: &VirtualMachine) -> PyResult { vm.call_special_method(obj, identifier!(vm, __getitem__), (needle,)) @@ -569,20 +583,20 @@ impl PyType { ); } _ if name == identifier!(ctx, __pow__) => { - toggle_subslot!(as_number, power, number_binary_op_wrapper!(__pow__)); + toggle_subslot!(as_number, power, number_ternary_op_wrapper!(__pow__)); } _ if name == identifier!(ctx, __rpow__) => { toggle_subslot!( as_number, right_power, - number_binary_right_op_wrapper!(__rpow__) + number_ternary_right_op_wrapper!(__rpow__) ); } _ if name == identifier!(ctx, __ipow__) => { toggle_subslot!( as_number, inplace_power, - number_binary_op_wrapper!(__ipow__) + number_ternary_op_wrapper!(__ipow__) ); } _ if name == identifier!(ctx, __lshift__) => { diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs index 7858b576f..65be5d415 100644 --- a/vm/src/vm/vm_ops.rs +++ b/vm/src/vm/vm_ops.rs @@ -2,7 +2,7 @@ use super::{PyMethod, VirtualMachine}; use crate::{ builtins::{PyInt, PyIntRef, PyStr, PyStrRef}, object::{AsObject, PyObject, PyObjectRef, PyResult}, - protocol::{PyIterReturn, PyNumberBinaryOp, PySequence}, + protocol::{PyIterReturn, PyNumberBinaryOp, PyNumberTernaryOp, PySequence}, types::PyComparisonOp, }; use num_traits::ToPrimitive; @@ -15,6 +15,14 @@ macro_rules! binary_func { }; } +macro_rules! ternary_func { + ($fn:ident, $op_slot:ident, $op:expr) => { + pub fn $fn(&self, a: &PyObject, b: &PyObject, c: &PyObject) -> PyResult { + self.ternary_op(a, b, c, PyNumberTernaryOp::$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 { @@ -29,6 +37,21 @@ macro_rules! inplace_binary_func { }; } +macro_rules! inplace_ternary_func { + ($fn:ident, $iop_slot:ident, $op_slot:ident, $op:expr) => { + pub fn $fn(&self, a: &PyObject, b: &PyObject, c: &PyObject) -> PyResult { + self.ternary_iop( + a, + b, + c, + PyNumberTernaryOp::$iop_slot, + PyNumberTernaryOp::$op_slot, + $op, + ) + } + }; +} + /// Collection of operators impl VirtualMachine { #[inline] @@ -229,10 +252,103 @@ impl VirtualMachine { Err(self.new_unsupported_binop_error(a, b, op)) } + fn ternary_op( + &self, + a: &PyObject, + b: &PyObject, + c: &PyObject, + op_slot: PyNumberTernaryOp, + op_str: &str, + ) -> PyResult { + let class_a = a.class(); + let class_b = b.class(); + let class_c = c.class(); + + let slot_a = class_a.slots.as_number.left_ternary_op(op_slot); + let mut slot_b = None; + + if !class_a.is(class_b) { + let slot_bb = class_b.slots.as_number.right_ternary_op(op_slot); + if slot_bb.map(|x| x as usize) != slot_a.map(|x| x as usize) { + slot_b = slot_bb; + } + } + + if let Some(slot_a) = slot_a { + if let Some(slot_bb) = slot_b { + if class_b.fast_issubclass(class_a) { + let ret = slot_bb(a, b, c, self)?; + if !ret.is(&self.ctx.not_implemented) { + return Ok(ret); + } + slot_b = None; + } + } + let ret = slot_a(a, b, c, self)?; + if !ret.is(&self.ctx.not_implemented) { + return Ok(ret); + } + } + + if let Some(slot_b) = slot_b { + let ret = slot_b(a, b, c, self)?; + if !ret.is(&self.ctx.not_implemented) { + return Ok(ret); + } + } + + if let Some(slot_c) = class_c.slots.as_number.left_ternary_op(op_slot) { + if slot_a.map_or(false, |slot_a| (slot_a as usize) != (slot_c as usize)) + && slot_b.map_or(false, |slot_b| (slot_b as usize) != (slot_c as usize)) + { + let ret = slot_c(a, b, c, self)?; + if !ret.is(&self.ctx.not_implemented) { + return Ok(ret); + } + } + } + + Err(if self.is_none(c) { + self.new_type_error(format!( + "unsupported operand type(s) for {}: \ + '{}' and '{}'", + op_str, + a.class(), + b.class() + )) + } else { + self.new_type_error(format!( + "unsupported operand type(s) for {}: \ + '{}' and '{}', '{}'", + op_str, + a.class(), + b.class(), + c.class() + )) + }) + } + + fn ternary_iop( + &self, + a: &PyObject, + b: &PyObject, + c: &PyObject, + iop_slot: PyNumberTernaryOp, + op_slot: PyNumberTernaryOp, + op_str: &str, + ) -> PyResult { + if let Some(slot) = a.class().slots.as_number.left_ternary_op(iop_slot) { + let x = slot(a, b, c, self)?; + if !x.is(&self.ctx.not_implemented) { + return Ok(x); + } + } + self.ternary_op(a, b, c, op_slot, op_str) + } + 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, "&"); @@ -244,7 +360,6 @@ impl VirtualMachine { 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, "&="); @@ -254,6 +369,9 @@ impl VirtualMachine { inplace_binary_func!(_itruediv, InplaceTrueDivide, TrueDivide, "/="); inplace_binary_func!(_imatmul, InplaceMatrixMultiply, MatrixMultiply, "@="); + ternary_func!(_pow, Power, "** or pow()"); + inplace_ternary_func!(_ipow, InplacePower, Power, "**="); + pub fn _add(&self, a: &PyObject, b: &PyObject) -> PyResult { let result = self.binary_op1(a, b, PyNumberBinaryOp::Add)?; if !result.is(&self.ctx.not_implemented) {