From e6db23f66d8741db0ffc526d8fd75373a5543e3e Mon Sep 17 00:00:00 2001
From: Dong-hee Na <donghee.na@python.org>
Date: Tue, 22 Aug 2023 06:22:18 +0900
Subject: [PATCH] gh-107265: Fix code_hash for ENTER_EXECUTOR case (#108188)

---
 Lib/test/test_capi/test_misc.py |  6 ++--
 Objects/codeobject.c            | 53 +++++++++++++++++++++------------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
index ea0504333ba..1cd4c56b49b 100644
--- a/Lib/test/test_capi/test_misc.py
+++ b/Lib/test/test_capi/test_misc.py
@@ -2341,7 +2341,7 @@ class TestOptimizerAPI(unittest.TestCase):
             long_loop()
             self.assertEqual(opt.get_count(), 10)
 
-    def test_code_richcompare(self):
+    def test_code_restore_for_ENTER_EXECUTOR(self):
         def testfunc(x):
             i = 0
             while i < x:
@@ -2350,7 +2350,9 @@ class TestOptimizerAPI(unittest.TestCase):
         opt = _testinternalcapi.get_counter_optimizer()
         with temporary_optimizer(opt):
             testfunc(1000)
-            self.assertEqual(testfunc.__code__, testfunc.__code__.replace())
+            code, replace_code  = testfunc.__code__, testfunc.__code__.replace()
+            self.assertEqual(code, replace_code)
+            self.assertEqual(hash(code), hash(replace_code))
 
 
 def get_first_executor(func):
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index c34905c3196..dca5804a91d 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -1781,30 +1781,33 @@ code_richcompare(PyObject *self, PyObject *other, int op)
     for (int i = 0; i < Py_SIZE(co); i++) {
         _Py_CODEUNIT co_instr = _PyCode_CODE(co)[i];
         _Py_CODEUNIT cp_instr = _PyCode_CODE(cp)[i];
+        uint8_t co_code = co_instr.op.code;
+        uint8_t co_arg = co_instr.op.arg;
+        uint8_t cp_code = cp_instr.op.code;
+        uint8_t cp_arg = cp_instr.op.arg;
 
-        if (co_instr.op.code == ENTER_EXECUTOR) {
-            const int exec_index = co_instr.op.arg;
+        if (co_code == ENTER_EXECUTOR) {
+            const int exec_index = co_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;
+            co_code = exec->vm_data.opcode;
+            co_arg = exec->vm_data.oparg;
         }
-        assert(co_instr.op.code != ENTER_EXECUTOR);
-        co_instr.op.code = _PyOpcode_Deopt[co_instr.op.code];
+        assert(co_code != ENTER_EXECUTOR);
+        co_code = _PyOpcode_Deopt[co_code];
 
-        if (cp_instr.op.code == ENTER_EXECUTOR) {
-            const int exec_index = cp_instr.op.arg;
+        if (cp_code == ENTER_EXECUTOR) {
+            const int exec_index = cp_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;
+            cp_code = exec->vm_data.opcode;
+            cp_arg = exec->vm_data.oparg;
         }
-        assert(cp_instr.op.code != ENTER_EXECUTOR);
-        cp_instr.op.code = _PyOpcode_Deopt[cp_instr.op.code];
+        assert(cp_code != ENTER_EXECUTOR);
+        cp_code = _PyOpcode_Deopt[cp_code];
 
-        eq = co_instr.cache == cp_instr.cache;
-        if (!eq) {
+        if (co_code != cp_code || co_arg != cp_arg) {
             goto unequal;
         }
-        i += _PyOpcode_Caches[co_instr.op.code];
+        i += _PyOpcode_Caches[co_code];
     }
 
     /* compare constants */
@@ -1883,10 +1886,22 @@ code_hash(PyCodeObject *co)
     SCRAMBLE_IN(co->co_firstlineno);
     SCRAMBLE_IN(Py_SIZE(co));
     for (int i = 0; i < Py_SIZE(co); i++) {
-        int deop = _Py_GetBaseOpcode(co, i);
-        SCRAMBLE_IN(deop);
-        SCRAMBLE_IN(_PyCode_CODE(co)[i].op.arg);
-        i += _PyOpcode_Caches[deop];
+        _Py_CODEUNIT co_instr = _PyCode_CODE(co)[i];
+        uint8_t co_code = co_instr.op.code;
+        uint8_t co_arg = co_instr.op.arg;
+        if (co_code == ENTER_EXECUTOR) {
+            _PyExecutorObject *exec = co->co_executors->executors[co_arg];
+            assert(exec != NULL);
+            assert(exec->vm_data.opcode != ENTER_EXECUTOR);
+            co_code = _PyOpcode_Deopt[exec->vm_data.opcode];
+            co_arg = exec->vm_data.oparg;
+        }
+        else {
+            co_code = _Py_GetBaseOpcode(co, i);
+        }
+        SCRAMBLE_IN(co_code);
+        SCRAMBLE_IN(co_arg);
+        i += _PyOpcode_Caches[co_code];
     }
     if ((Py_hash_t)uhash == -1) {
         return -2;