Binary ops with number protocol (partially imported)

Karatuss · youknowone · commit e5c0961ee429 · 2023-03-03T12:26:20.000+09:00
from #4139 #4139
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, PyNumberMethods, PyNumberMethodsOffset};
 pub use sequence::{PySequence, PySequenceMethods};
diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs
@@ -8,6 +8,7 @@ use crate::{
     VirtualMachine,
 };
 use crossbeam_utils::atomic::AtomicCell;
+use std::ptr;
 
 type UnaryFunc<R = PyObjectRef> = AtomicCell<Option<fn(PyNumber, &VirtualMachine) -> PyResult<R>>>;
 type BinaryFunc<R = PyObjectRef> =
@@ -109,6 +110,7 @@ impl PyObject {
 }
 
 #[derive(Default)]
+// #[repr(C)]
 pub struct PyNumberMethods {
     /* Number implementations must check *both*
     arguments for proper type and implement the necessary conversions
@@ -199,6 +201,46 @@ impl PyNumberMethods {
     };
 }
 
+pub enum PyNumberMethodsOffset {
+    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,
+    InplaceAnd,
+    InplaceXor,
+    InplaceOr,
+    FloorDivide,
+    TrueDivide,
+    InplaceFloorDivide,
+    InplaceTrueDivide,
+    Index,
+    MatrixMultiply,
+    InplaceMatrixMultiply,
+}
+
+#[derive(Copy, Clone)]
 pub struct PyNumber<'a> {
     pub obj: &'a PyObject,
     methods: &'a PyNumberMethods,
@@ -220,8 +262,12 @@ impl PyNumber<'_> {
         obj.class().mro_find_map(|x| x.slots.as_number.load())
     }
 
-    pub fn methods(&self) -> &PyNumberMethods {
-        self.methods
+    pub fn methods<'a>(
+        &'a self,
+        op_slot: &'a PyNumberMethodsOffset,
+        vm: &VirtualMachine,
+    ) -> PyResult<&BinaryFunc> {
+        op_slot.method(self.methods, vm)
     }
 
     // PyNumber_Check
@@ -238,12 +284,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<PyResult<PyIntRef>> {
-        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(
@@ -267,7 +313,7 @@ impl PyNumber<'_> {
 
     #[inline]
     pub fn index(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
-        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(
@@ -291,7 +337,7 @@ impl PyNumber<'_> {
 
     #[inline]
     pub fn float(self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
-        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
@@ -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, PyNumberMethodsOffset},
         py_io,
         readline::{Readline, ReadlineResult},
         stdlib::sys,
@@ -610,13 +610,9 @@ 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, PyNumberMethodsOffset::Power, "pow", |vm, _, _| {
+                Ok(vm.ctx.not_implemented())
+            }),
             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
