Improve: binaryops with Number Protocol

xiaozhiyan · xiaozhiyan · commit 7c7f4debf00a · 2023-03-02T23:53:34.000+09:00
diff --git a/vm/src/protocol/mod.rs b/vm/src/protocol/mod.rs
@@ -8,5 +8,5 @@ mod sequence;
 pub use buffer::{BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, VecBuffer};
 pub use iter::{PyIter, PyIterIter, PyIterReturn};
 pub use mapping::{PyMapping, PyMappingMethods};
-pub use number::{PyNumber, PyNumberMethods};
+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
@@ -136,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,
@@ -182,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),
@@ -197,8 +195,73 @@ 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.inplace_floor_divide,
+            TrueDivide => &self.inplace_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 PyNumberBinaryOpSlot {
+    Add,
+    Subtract,
+    Multiply,
+    Remainder,
+    Divmod,
+    Power,
+    Lshift,
+    Rshift,
+    And,
+    Xor,
+    Or,
+    InplaceAdd,
+    InplaceSubtract,
+    InplaceMultiply,
+    InplaceRemainder,
+    InplacePower,
+    InplaceLshift,
+    InplaceRshift,
+    InplaceAnd,
+    InplaceXor,
+    InplaceOr,
+    FloorDivide,
+    TrueDivide,
+    InplaceFloorDivide,
+    InplaceTrueDivide,
+    MatrixMultiply,
+    InplaceMatrixMultiply,
+}
+
+#[derive(Copy, Clone)]
 pub struct PyNumber<'a> {
     pub obj: &'a PyObject,
     methods: &'a PyNumberMethods,
@@ -224,6 +287,10 @@ impl PyNumber<'_> {
         self.methods
     }
 
+    pub fn get_binary_op(&self, op_slot: &PyNumberBinaryOpSlot) -> PyResult<&BinaryFunc> {
+        self.methods().get_binary_op(op_slot)
+    }
+
     // PyNumber_Check
     pub fn check(obj: &PyObject) -> bool {
         let Some(methods) = Self::find_methods(obj) else {
diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs
@@ -22,7 +22,7 @@ mod builtins {
             ArgBytesLike, ArgCallable, ArgIntoBool, ArgIterable, ArgMapping, ArgStrOrBytesLike,
             Either, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs, PyArithmeticValue,
         },
-        protocol::{PyIter, PyIterReturn},
+        protocol::{PyIter, PyIterReturn, PyNumberBinaryOpSlot},
         py_io,
         readline::{Readline, ReadlineResult},
         stdlib::sys,
@@ -610,13 +610,7 @@ mod builtins {
             modulus,
         } = args;
         match modulus {
-            None => vm.call_or_reflection(
-                &x,
-                &y,
-                identifier!(vm, __pow__),
-                identifier!(vm, __rpow__),
-                |vm, x, y| Err(vm.new_unsupported_binop_error(x, y, "pow")),
-            ),
+            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
