Fix power, inplace_power as ternary function

Signed-off-by: snowapril <sinjihng@gmail.com>

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};

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<PyNumberBinaryFunc>,
     pub remainder: Option<PyNumberBinaryFunc>,
     pub divmod: Option<PyNumberBinaryFunc>,
-    pub power: Option<PyNumberBinaryFunc>,
+    pub power: Option<PyNumberTernaryFunc>,
     pub negative: Option<PyNumberUnaryFunc>,
     pub positive: Option<PyNumberUnaryFunc>,
     pub absolute: Option<PyNumberUnaryFunc>,
@@ -139,7 +140,7 @@ pub struct PyNumberMethods {
     pub inplace_subtract: Option<PyNumberBinaryFunc>,
     pub inplace_multiply: Option<PyNumberBinaryFunc>,
     pub inplace_remainder: Option<PyNumberBinaryFunc>,
-    pub inplace_power: Option<PyNumberBinaryFunc>,
+    pub inplace_power: Option<PyNumberTernaryFunc>,
     pub inplace_lshift: Option<PyNumberBinaryFunc>,
     pub inplace_rshift: Option<PyNumberBinaryFunc>,
     pub inplace_and: Option<PyNumberBinaryFunc>,
@@ -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<Option<PyNumberBinaryFunc>>,
@@ -241,7 +246,7 @@ pub struct PyNumberSlots {
     pub multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub divmod: AtomicCell<Option<PyNumberBinaryFunc>>,
-    pub power: AtomicCell<Option<PyNumberBinaryFunc>>,
+    pub power: AtomicCell<Option<PyNumberTernaryFunc>>,
     pub negative: AtomicCell<Option<PyNumberUnaryFunc>>,
     pub positive: AtomicCell<Option<PyNumberUnaryFunc>>,
     pub absolute: AtomicCell<Option<PyNumberUnaryFunc>>,
@@ -260,7 +265,7 @@ pub struct PyNumberSlots {
     pub right_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub right_remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub right_divmod: AtomicCell<Option<PyNumberBinaryFunc>>,
-    pub right_power: AtomicCell<Option<PyNumberBinaryFunc>>,
+    pub right_power: AtomicCell<Option<PyNumberTernaryFunc>>,
     pub right_lshift: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub right_rshift: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub right_and: AtomicCell<Option<PyNumberBinaryFunc>>,
@@ -271,7 +276,7 @@ pub struct PyNumberSlots {
     pub inplace_subtract: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub inplace_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub inplace_remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
-    pub inplace_power: AtomicCell<Option<PyNumberBinaryFunc>>,
+    pub inplace_power: AtomicCell<Option<PyNumberTernaryFunc>>,
     pub inplace_lshift: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub inplace_rshift: AtomicCell<Option<PyNumberBinaryFunc>>,
     pub inplace_and: AtomicCell<Option<PyNumberBinaryFunc>>,
@@ -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<PyNumberTernaryFunc> {
+        use PyNumberTernaryOp::*;
+        match op_slot {
+            Power => self.power.load(),
+            InplacePower => self.inplace_power.load(),
+        }
+    }
+
+    pub fn right_ternary_op(&self, op_slot: PyNumberTernaryOp) -> Option<PyNumberTernaryFunc> {
+        use PyNumberTernaryOp::*;
+        match op_slot {
+            Power => self.right_power.load(),
+            _ => None,
+        }
+    }
 }
 #[derive(Copy, Clone)]
 pub struct PyNumber<'a>(&'a PyObject);

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<K: ToPyObject>(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__) => {

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) {