gh-91625: Don't ignore extended args of adaptive opcodes (GH-91626)

5 files changed, 66 insertions, 18 deletions

diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
index 5d36cb9..378ff52 100644
--- a/Lib/test/test_descr.py
+++ b/Lib/test/test_descr.py
@@ -1961,6 +1961,20 @@ order (MRO) for bases """
         del a[0:10]
         self.assertEqual(a.delitem, (slice(0, 10)))
 
+    def test_load_attr_extended_arg(self):
+        # https://github.com/python/cpython/issues/91625
+        class Numbers:
+            def __getattr__(self, attr):
+                return int(attr.lstrip("_"))
+        attrs = ", ".join(f"Z._{n:03d}" for n in range(280))
+        code = f"def number_attrs(Z):\n    return [ {attrs} ]"
+        ns = {}
+        exec(code, ns)
+        number_attrs = ns["number_attrs"]
+        # Warm up the the function for quickening (PEP 659)
+        for _ in range(30):
+            self.assertEqual(number_attrs(Numbers()), list(range(280)))
+
     def test_methods(self):
         # Testing methods...
         class C(object):
@@ -4435,8 +4449,8 @@ order (MRO) for bases """
                 raise RuntimeError(f"Premature access to sys.stdout.{attr}")
 
         with redirect_stdout(StdoutGuard()):
-             with self.assertRaises(RuntimeError):
-                 print("Oops!")
+            with self.assertRaises(RuntimeError):
+                print("Oops!")
 
     def test_vicious_descriptor_nonsense(self):
         # Testing vicious_descriptor_nonsense...
diff --git a/Lib/test/test_dynamic.py b/Lib/test/test_dynamic.py
index 3ae090f..3e0fcf4 100644
--- a/Lib/test/test_dynamic.py
+++ b/Lib/test/test_dynamic.py
@@ -133,6 +133,18 @@ class RebindBuiltinsTests(unittest.TestCase):
 
         self.assertEqual(foo(), 7)
 
+    def test_load_global_specialization_failure_keeps_oparg(self):
+        # https://github.com/python/cpython/issues/91625
+        class MyGlobals(dict):
+            def __missing__(self, key):
+                return int(key.removeprefix("_number_"))
+
+        code = "lambda: " + "+".join(f"_number_{i}" for i in range(1000))
+        sum_1000 = eval(code, MyGlobals())
+        expected = sum(range(1000))
+        # Warm up the the function for quickening (PEP 659)
+        for _ in range(30):
+            self.assertEqual(sum_1000(), expected)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_unpack.py b/Lib/test/test_unpack.py
index 472c834..f5ca1d4 100644
--- a/Lib/test/test_unpack.py
+++ b/Lib/test/test_unpack.py
@@ -151,5 +151,21 @@ def load_tests(loader, tests, pattern):
     return tests
 
 
+class TestCornerCases(unittest.TestCase):
+    def test_extended_oparg_not_ignored(self):
+        # https://github.com/python/cpython/issues/91625
+        target = "(" + "y,"*400 + ")"
+        code = f"""def unpack_400(x):
+            {target} = x
+            return y
+        """
+        ns = {}
+        exec(code, ns)
+        unpack_400 = ns["unpack_400"]
+        # Warm up the the function for quickening (PEP 659)
+        for _ in range(30):
+            y = unpack_400(range(400))
+            self.assertEqual(y, 399)
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-04-17-02-55-38.gh-issue-91625.80CrC7.rst b/Misc/NEWS.d/next/Core and Builtins/2022-04-17-02-55-38.gh-issue-91625.80CrC7.rst
new file mode 100644
index 0000000..ea5b57b
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-04-17-02-55-38.gh-issue-91625.80CrC7.rst
@@ -0,0 +1 @@
+Fixed a bug in which adaptive opcodes ignored any preceding ``EXTENDED_ARG``\ s on specialization failure.
diff --git a/Python/ceval.c b/Python/ceval.c
index d358a31..f523e52 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1351,6 +1351,13 @@ eval_frame_handle_pending(PyThreadState *tstate)
         DISPATCH_GOTO(); \
     }
 
+#define NOTRACE_DISPATCH_SAME_OPARG() \
+    { \
+        opcode = _Py_OPCODE(*next_instr); \
+        PRE_DISPATCH_GOTO(); \
+        DISPATCH_GOTO(); \
+    }
+
 #define CHECK_EVAL_BREAKER() \
     _Py_CHECK_EMSCRIPTEN_SIGNALS_PERIODICALLY(); \
     if (_Py_atomic_load_relaxed(eval_breaker)) { \
@@ -2158,7 +2165,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_BinarySubscr(container, sub, next_instr) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(BINARY_SUBSCR, deferred);
@@ -2323,7 +2330,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_StoreSubscr(container, sub, next_instr) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(STORE_SUBSCR, deferred);
@@ -2813,7 +2820,7 @@ handle_eval_breaker:
                 PyObject *seq = TOP();
                 next_instr--;
                 _Py_Specialize_UnpackSequence(seq, next_instr, oparg);
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(UNPACK_SEQUENCE, deferred);
@@ -3056,7 +3063,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_LoadGlobal(GLOBALS(), BUILTINS(), next_instr, name) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(LOAD_GLOBAL, deferred);
@@ -3481,7 +3488,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_LoadAttr(owner, next_instr, name) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(LOAD_ATTR, deferred);
@@ -3590,7 +3597,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_StoreAttr(owner, next_instr, name) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(STORE_ATTR, deferred);
@@ -3718,7 +3725,7 @@ handle_eval_breaker:
                 PyObject *left = SECOND();
                 next_instr--;
                 _Py_Specialize_CompareOp(left, right, next_instr, oparg);
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(COMPARE_OP, deferred);
@@ -4523,7 +4530,7 @@ handle_eval_breaker:
                 if (_Py_Specialize_LoadMethod(owner, next_instr, name) < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(LOAD_METHOD, deferred);
@@ -4797,7 +4804,7 @@ handle_eval_breaker:
                 if (err < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(PRECALL, deferred);
@@ -4818,7 +4825,7 @@ handle_eval_breaker:
                 if (err < 0) {
                     goto error;
                 }
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(CALL, deferred);
@@ -5545,7 +5552,7 @@ handle_eval_breaker:
                 PyObject *rhs = TOP();
                 next_instr--;
                 _Py_Specialize_BinaryOp(lhs, rhs, next_instr, oparg, &GETLOCAL(0));
-                DISPATCH();
+                NOTRACE_DISPATCH_SAME_OPARG();
             }
             else {
                 STAT_INC(BINARY_OP, deferred);
@@ -5564,11 +5571,9 @@ handle_eval_breaker:
 
         TARGET(EXTENDED_ARG) {
             assert(oparg);
-            int oldoparg = oparg;
-            NEXTOPARG();
-            oparg |= oldoparg << 8;
-            PRE_DISPATCH_GOTO();
-            DISPATCH_GOTO();
+            oparg <<= 8;
+            oparg |= _Py_OPARG(*next_instr);
+            NOTRACE_DISPATCH_SAME_OPARG();
         }
 
         TARGET(CACHE) {