gh-98831: rewrite BEFORE_ASYNC_WITH and END_ASYNC_FOR in the instruction definition DSL (#101458)

3 files changed, 40 insertions, 38 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 5b3bf4e..336088e 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -770,18 +770,16 @@ dummy_func(
             ERROR_IF(val == NULL, error);
         }
 
-        // stack effect: (__0, __1 -- )
-        inst(END_ASYNC_FOR) {
-            PyObject *val = POP();
-            assert(val && PyExceptionInstance_Check(val));
-            if (PyErr_GivenExceptionMatches(val, PyExc_StopAsyncIteration)) {
-                Py_DECREF(val);
-                Py_DECREF(POP());
+        inst(END_ASYNC_FOR, (awaitable, exc -- )) {
+            assert(exc && PyExceptionInstance_Check(exc));
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration)) {
+                DECREF_INPUTS();
             }
             else {
-                PyObject *exc = Py_NewRef(PyExceptionInstance_Class(val));
-                PyObject *tb = PyException_GetTraceback(val);
-                _PyErr_Restore(tstate, exc, val, tb);
+                Py_INCREF(exc);
+                PyObject *typ = Py_NewRef(PyExceptionInstance_Class(exc));
+                PyObject *tb = PyException_GetTraceback(exc);
+                _PyErr_Restore(tstate, typ, exc, tb);
                 goto exception_unwind;
             }
         }
@@ -2266,10 +2264,7 @@ dummy_func(
             DISPATCH_INLINED(gen_frame);
         }
 
-        // stack effect: ( -- __0)
-        inst(BEFORE_ASYNC_WITH) {
-            PyObject *mgr = TOP();
-            PyObject *res;
+        inst(BEFORE_ASYNC_WITH, (mgr -- exit, res)) {
             PyObject *enter = _PyObject_LookupSpecial(mgr, &_Py_ID(__aenter__));
             if (enter == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
@@ -2280,7 +2275,7 @@ dummy_func(
                 }
                 goto error;
             }
-            PyObject *exit = _PyObject_LookupSpecial(mgr, &_Py_ID(__aexit__));
+            exit = _PyObject_LookupSpecial(mgr, &_Py_ID(__aexit__));
             if (exit == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
                     _PyErr_Format(tstate, PyExc_TypeError,
@@ -2292,13 +2287,13 @@ dummy_func(
                 Py_DECREF(enter);
                 goto error;
             }
-            SET_TOP(exit);
-            Py_DECREF(mgr);
+            DECREF_INPUTS();
             res = _PyObject_CallNoArgs(enter);
             Py_DECREF(enter);
-            if (res == NULL)
-                goto error;
-            PUSH(res);
+            if (res == NULL) {
+                Py_DECREF(exit);
+                ERROR_IF(true, error);
+            }
             PREDICT(GET_AWAITABLE);
         }
 
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 661fe27..d70d64e 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -981,18 +981,21 @@
         }
 
         TARGET(END_ASYNC_FOR) {
-            PyObject *val = POP();
-            assert(val && PyExceptionInstance_Check(val));
-            if (PyErr_GivenExceptionMatches(val, PyExc_StopAsyncIteration)) {
-                Py_DECREF(val);
-                Py_DECREF(POP());
+            PyObject *exc = PEEK(1);
+            PyObject *awaitable = PEEK(2);
+            assert(exc && PyExceptionInstance_Check(exc));
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration)) {
+                Py_DECREF(awaitable);
+                Py_DECREF(exc);
             }
             else {
-                PyObject *exc = Py_NewRef(PyExceptionInstance_Class(val));
-                PyObject *tb = PyException_GetTraceback(val);
-                _PyErr_Restore(tstate, exc, val, tb);
+                Py_INCREF(exc);
+                PyObject *typ = Py_NewRef(PyExceptionInstance_Class(exc));
+                PyObject *tb = PyException_GetTraceback(exc);
+                _PyErr_Restore(tstate, typ, exc, tb);
                 goto exception_unwind;
             }
+            STACK_SHRINK(2);
             DISPATCH();
         }
 
@@ -2699,7 +2702,8 @@
         }
 
         TARGET(BEFORE_ASYNC_WITH) {
-            PyObject *mgr = TOP();
+            PyObject *mgr = PEEK(1);
+            PyObject *exit;
             PyObject *res;
             PyObject *enter = _PyObject_LookupSpecial(mgr, &_Py_ID(__aenter__));
             if (enter == NULL) {
@@ -2711,7 +2715,7 @@
                 }
                 goto error;
             }
-            PyObject *exit = _PyObject_LookupSpecial(mgr, &_Py_ID(__aexit__));
+            exit = _PyObject_LookupSpecial(mgr, &_Py_ID(__aexit__));
             if (exit == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
                     _PyErr_Format(tstate, PyExc_TypeError,
@@ -2723,13 +2727,16 @@
                 Py_DECREF(enter);
                 goto error;
             }
-            SET_TOP(exit);
             Py_DECREF(mgr);
             res = _PyObject_CallNoArgs(enter);
             Py_DECREF(enter);
-            if (res == NULL)
-                goto error;
-            PUSH(res);
+            if (res == NULL) {
+                Py_DECREF(exit);
+                if (true) goto pop_1_error;
+            }
+            STACK_GROW(1);
+            POKE(1, res);
+            POKE(2, exit);
             PREDICT(GET_AWAITABLE);
             DISPATCH();
         }
diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
index c40e40f..171fed3 100644
--- a/Python/opcode_metadata.h
+++ b/Python/opcode_metadata.h
@@ -109,7 +109,7 @@ _PyOpcode_num_popped(int opcode, int oparg) {
         case PREP_RERAISE_STAR:
             return 2;
         case END_ASYNC_FOR:
-            return -1;
+            return 2;
         case CLEANUP_THROW:
             return -1;
         case LOAD_ASSERTION_ERROR:
@@ -271,7 +271,7 @@ _PyOpcode_num_popped(int opcode, int oparg) {
         case FOR_ITER_GEN:
             return -1;
         case BEFORE_ASYNC_WITH:
-            return -1;
+            return 1;
         case BEFORE_WITH:
             return 1;
         case WITH_EXCEPT_START:
@@ -455,7 +455,7 @@ _PyOpcode_num_pushed(int opcode, int oparg) {
         case PREP_RERAISE_STAR:
             return 1;
         case END_ASYNC_FOR:
-            return -1;
+            return 0;
         case CLEANUP_THROW:
             return -1;
         case LOAD_ASSERTION_ERROR:
@@ -617,7 +617,7 @@ _PyOpcode_num_pushed(int opcode, int oparg) {
         case FOR_ITER_GEN:
             return -1;
         case BEFORE_ASYNC_WITH:
-            return -1;
+            return 2;
         case BEFORE_WITH:
             return 2;
         case WITH_EXCEPT_START: