Add test and compare oparg too

corona10 · corona10 · commit fb0ab4e4f329 · 2023-08-21T12:22:52.000+09:00
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
@@ -2341,6 +2341,17 @@ def long_loop():
             long_loop()
             self.assertEqual(opt.get_count(), 10)
 
+    def test_code_richcompare(self):
+        def testfunc(x):
+            i = 0
+            while i < x:
+                i += 1
+
+        opt = _testinternalcapi.get_counter_optimizer()
+        with temporary_optimizer(opt):
+            testfunc(1000)
+            self.assertEqual(testfunc.__code__, testfunc.__code__.replace())
+
 
 def get_first_executor(func):
     code = func.__code__
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
@@ -1786,6 +1786,7 @@ code_richcompare(PyObject *self, PyObject *other, int op)
             const int exec_index = co_instr.op.arg;
             _PyExecutorObject *exec = co->co_executors->executors[exec_index];
             co_instr.op.code = exec->vm_data.opcode;
+            co_instr.op.arg = exec->vm_data.oparg;
         }
         assert(co_instr.op.code != ENTER_EXECUTOR);
         co_instr.op.code = _PyOpcode_Deopt[co_instr.op.code];
@@ -1794,6 +1795,7 @@ code_richcompare(PyObject *self, PyObject *other, int op)
             const int exec_index = cp_instr.op.arg;
             _PyExecutorObject *exec = cp->co_executors->executors[exec_index];
             cp_instr.op.code = exec->vm_data.opcode;
+            cp_instr.op.arg = exec->vm_data.oparg;
         }
         assert(cp_instr.op.code != ENTER_EXECUTOR);
         cp_instr.op.code = _PyOpcode_Deopt[cp_instr.op.code];
