From 1ed18c012af12d0c9dd990ade8d42eb9dee74e56 Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Sun, 22 May 2022 21:00:18 +0200
Subject: [PATCH] fix vm.to_index now use number protocol

---
 Lib/test/test_index.py    |  2 --
 vm/src/protocol/number.rs | 32 +++++++++++++++++++++++---------
 vm/src/vm/vm_ops.rs       | 24 +++---------------------
 3 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/Lib/test/test_index.py b/Lib/test/test_index.py
index 1fac13259..cbdc56c80 100644
--- a/Lib/test/test_index.py
+++ b/Lib/test/test_index.py
@@ -71,8 +71,6 @@ class BaseTestCase(unittest.TestCase):
         self.assertIs(type(direct_index), int)
         #self.assertIs(type(operator_index), int)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_index_returns_int_subclass(self):
         class BadInt:
             def __index__(self):
diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs
index 68b575c55..2868ca8bd 100644
--- a/vm/src/protocol/number.rs
+++ b/vm/src/protocol/number.rs
@@ -134,9 +134,11 @@ impl PyNumber<'_> {
                     ),
                     1,
                     vm,
-                )?
+                )?;
+                Ok(vm.ctx.new_int(ret.as_bigint().clone()))
+            } else {
+                Ok(ret)
             }
-            Ok(ret)
         } else if self.methods(vm).index.is_some() {
             self.index(vm)
         } else if let Ok(Ok(f)) =
@@ -173,9 +175,13 @@ impl PyNumber<'_> {
         }
     }
 
-    pub fn index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
+    pub fn index_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyIntRef>> {
         if self.obj.class().is(PyInt::class(vm)) {
-            Ok(unsafe { self.obj.to_owned().downcast_unchecked::<PyInt>() })
+            Ok(Some(unsafe {
+                self.obj.to_owned().downcast_unchecked::<PyInt>()
+            }))
+        } else if let Some(i) = self.obj.downcast_ref::<PyInt>() {
+            Ok(Some(i.to_owned()))
         } else if let Some(f) = self.methods(vm).index {
             let ret = f(self, vm)?;
             if !ret.class().is(PyInt::class(vm)) {
@@ -189,15 +195,23 @@ impl PyNumber<'_> {
                     ),
                     1,
                     vm,
-                )?
+                )?;
+                Ok(Some(vm.ctx.new_int(ret.as_bigint().clone())))
+            } else {
+                Ok(Some(ret))
             }
-            Ok(ret)
         } else {
-            Err(vm.new_type_error(format!(
+            Ok(None)
+        }
+    }
+
+    pub fn index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
+        self.index_opt(vm)?.ok_or_else(|| {
+            vm.new_type_error(format!(
                 "'{}' object cannot be interpreted as an integer",
                 self.obj.class()
-            )))
-        }
+            ))
+        })
     }
 
     pub fn float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs
index e0606aefe..1c325595c 100644
--- a/vm/src/vm/vm_ops.rs
+++ b/vm/src/vm/vm_ops.rs
@@ -3,36 +3,18 @@ use crate::{
     builtins::{PyInt, PyIntRef, PyStrInterned},
     function::PyArithmeticValue,
     object::{AsObject, PyObject, PyObjectRef, PyResult},
-    protocol::PyIterReturn,
+    protocol::{PyIterReturn, PyNumber},
     types::PyComparisonOp,
 };
 
 /// Collection of operators
 impl VirtualMachine {
     pub fn to_index_opt(&self, obj: PyObjectRef) -> Option<PyResult<PyIntRef>> {
-        match obj.downcast() {
-            Ok(val) => Some(Ok(val)),
-            Err(obj) => self
-                .get_method(obj, identifier!(self, __index__))
-                .map(|index| {
-                    // TODO: returning strict subclasses of int in __index__ is deprecated
-                    self.invoke(&index?, ())?.downcast().map_err(|bad| {
-                        self.new_type_error(format!(
-                            "__index__ returned non-int (type {})",
-                            bad.class().name()
-                        ))
-                    })
-                }),
-        }
+        PyNumber::from(obj.as_ref()).index_opt(self).transpose()
     }
 
     pub fn to_index(&self, obj: &PyObject) -> PyResult<PyIntRef> {
-        self.to_index_opt(obj.to_owned()).unwrap_or_else(|| {
-            Err(self.new_type_error(format!(
-                "'{}' object cannot be interpreted as an integer",
-                obj.class().name()
-            )))
-        })
+        PyNumber::from(obj).index(self)
     }
 
     #[inline]