Upgrade with number protocol

Karatuss · Karatuss · commit 6b4373aa1f2d · 2022-08-28T16:12:02.000+09:00
diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs
@@ -611,8 +611,24 @@ impl AsSequence for PyBytes {
 impl AsNumber for PyBytes {
     fn as_number() -> &'static PyNumberMethods {
         static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            remainder: atomic_func!(|number, other, vm| {
-                PyBytes::number_downcast(number)
+            // ! `multiply` MUST be removed from number protocol later!!
+            // multiply: atomic_func!(|num, other, vm| {
+            //     let bytes_num = PyBytes::number_downcast(num);
+            //     PyBytes::mul(
+            //         if let Some(target) = num.obj.downcast_ref::<PyBytes>() {
+            //             target.into_ref(vm)
+            //         } else {
+
+            //         },
+            //         other.to_number().int(vm).map_or(
+            //             1 as isize,
+            //             |i| i.unwrap_or_else(|| PyInt::from(1).into_ref(vm)).as_u32_mask() as isize
+            //         ),
+            //         vm
+            //     ).to_pyresult(vm)
+            // }),
+            remainder: atomic_func!(|num, other, vm| {
+                PyBytes::number_downcast(num)
                     .mod_(other.to_owned(), vm)
                     .to_pyresult(vm)
             }),
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
@@ -216,6 +216,10 @@ impl PyNumberMethods {
     };
 }
 
+pub enum PyNumberMethodsOffset {
+    Multiply = 0,
+}
+
 pub struct PyNumber<'a> {
     pub obj: &'a PyObject,
     methods: &'a PyNumberMethods,
diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs
@@ -3,7 +3,7 @@ use crate::{
     builtins::{PyInt, PyIntRef, PyStrInterned},
     function::PyArithmeticValue,
     object::{AsObject, PyObject, PyObjectRef, PyResult},
-    protocol::PyIterReturn,
+    protocol::{PyIterReturn, PyNumberMethodsOffset},
     types::PyComparisonOp,
 };
 
@@ -174,6 +174,62 @@ impl VirtualMachine {
         })
     }
 
+    #[allow(unused_variables)]
+    fn binary_op1(&self, a: &PyObject, b: &PyObject, op_slot: PyNumberMethodsOffset) -> PyResult {
+        let num_a = a.to_number();
+        let num_b = b.to_number();
+
+        // ! `multiply` should be changed after this function works normally!!
+        let slot_a = num_a.methods().multiply.load();
+        let slot_b = num_b.methods().multiply.load();
+
+        if let Some(slot_a) = slot_a {
+            if let Some(slot_b) = slot_b {
+                // Check if a is subclass of b
+                if a.fast_isinstance(&b.class()) {
+                    let ret = slot_b(&num_b, a, self)?;
+                    if ret.rich_compare_bool(
+                        self.ctx.not_implemented.as_object(),
+                        PyComparisonOp::Ne,
+                        self,
+                    )? {
+                        return Ok(ret);
+                    }
+                }
+            }
+            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);
+            }
+        }
+        // No slot_a
+        if let Some(slot_b) = slot_b {
+            let ret = slot_b(&num_b, a, self)?;
+            if ret.rich_compare_bool(
+                self.ctx.not_implemented.as_object(),
+                PyComparisonOp::Ne,
+                self,
+            )? {
+                return Ok(ret);
+            }
+        }
+
+        // Both slot_a & slot_b don't exist.
+        // Ok(self.ctx.not_implemented())
+
+        self.call_or_reflection(
+            a,
+            b,
+            identifier!(self, __mul__),
+            identifier!(self, __rmul__),
+            |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "*")),
+        )
+    }
+
     pub fn _sub(&self, a: &PyObject, b: &PyObject) -> PyResult {
         self.call_or_reflection(
             a,
@@ -219,13 +275,15 @@ impl VirtualMachine {
     }
 
     pub fn _mul(&self, a: &PyObject, b: &PyObject) -> PyResult {
-        self.call_or_reflection(
-            a,
-            b,
-            identifier!(self, __mul__),
-            identifier!(self, __rmul__),
-            |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "*")),
-        )
+        self.binary_op1(a, b, PyNumberMethodsOffset::Multiply)
+
+        // self.call_or_reflection(
+        //     a,
+        //     b,
+        //     identifier!(self, __mul__),
+        //     identifier!(self, __rmul__),
+        //     |vm, a, b| Err(vm.new_unsupported_binop_error(a, b, "*")),
+        // )
     }
 
     pub fn _imul(&self, a: &PyObject, b: &PyObject) -> PyResult {
