From 8c183cddd3d0031c6d397738f73c20ee6bf61ce8 Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Tue, 24 Jan 2023 09:43:16 +0000
Subject: [PATCH] gh-98831: rewrite CHECK_EG_MATCH opcode in the instruction
 definition DSL (#101269)

---
 Python/bytecodes.c         | 40 ++++++++++---------------------------
 Python/ceval.c             |  2 +-
 Python/generated_cases.c.h | 41 ++++++++++++++------------------------
 Python/opcode_metadata.h   |  2 +-
 4 files changed, 27 insertions(+), 58 deletions(-)

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 6088fa45ac6..47bbe1a99e3 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -1894,44 +1894,24 @@ dummy_func(
             b = Py_NewRef((res^oparg) ? Py_True : Py_False);
         }
 
-        // stack effect: ( -- )
-        inst(CHECK_EG_MATCH) {
-            PyObject *match_type = POP();
+        inst(CHECK_EG_MATCH, (exc_value, match_type -- rest, match)) {
             if (check_except_star_type_valid(tstate, match_type) < 0) {
-                Py_DECREF(match_type);
-                goto error;
+                DECREF_INPUTS();
+                ERROR_IF(true, error);
             }
 
-            PyObject *exc_value = TOP();
-            PyObject *match = NULL, *rest = NULL;
+            match = NULL;
+            rest = NULL;
             int res = exception_group_match(exc_value, match_type,
                                             &match, &rest);
-            Py_DECREF(match_type);
-            if (res < 0) {
-                goto error;
-            }
+            DECREF_INPUTS();
+            ERROR_IF(res < 0, error);
 
-            if (match == NULL || rest == NULL) {
-                assert(match == NULL);
-                assert(rest == NULL);
-                goto error;
-            }
-            if (Py_IsNone(match)) {
-                PUSH(match);
-                Py_XDECREF(rest);
-            }
-            else {
-                /* Total or partial match - update the stack from
-                 * [val]
-                 * to
-                 * [rest, match]
-                 * (rest can be Py_None)
-                 */
+            assert((match == NULL) == (rest == NULL));
+            ERROR_IF(match == NULL, error);
 
-                SET_TOP(rest);
-                PUSH(match);
+            if (!Py_IsNone(match)) {
                 PyErr_SetExcInfo(NULL, Py_NewRef(match), NULL);
-                Py_DECREF(exc_value);
             }
         }
 
diff --git a/Python/ceval.c b/Python/ceval.c
index a97313c773e..95eb99b4533 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1946,7 +1946,7 @@ exception_group_match(PyObject* exc_value, PyObject *match_type,
     }
     /* no match */
     *match = Py_NewRef(Py_None);
-    *rest = Py_NewRef(Py_None);
+    *rest = Py_NewRef(exc_value);
     return 0;
 }
 
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 46d421f98ac..c1eb4000883 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -2251,43 +2251,32 @@
         }
 
         TARGET(CHECK_EG_MATCH) {
-            PyObject *match_type = POP();
+            PyObject *match_type = PEEK(1);
+            PyObject *exc_value = PEEK(2);
+            PyObject *rest;
+            PyObject *match;
             if (check_except_star_type_valid(tstate, match_type) < 0) {
+                Py_DECREF(exc_value);
                 Py_DECREF(match_type);
-                goto error;
+                if (true) goto pop_2_error;
             }
 
-            PyObject *exc_value = TOP();
-            PyObject *match = NULL, *rest = NULL;
+            match = NULL;
+            rest = NULL;
             int res = exception_group_match(exc_value, match_type,
                                             &match, &rest);
+            Py_DECREF(exc_value);
             Py_DECREF(match_type);
-            if (res < 0) {
-                goto error;
-            }
+            if (res < 0) goto pop_2_error;
 
-            if (match == NULL || rest == NULL) {
-                assert(match == NULL);
-                assert(rest == NULL);
-                goto error;
-            }
-            if (Py_IsNone(match)) {
-                PUSH(match);
-                Py_XDECREF(rest);
-            }
-            else {
-                /* Total or partial match - update the stack from
-                 * [val]
-                 * to
-                 * [rest, match]
-                 * (rest can be Py_None)
-                 */
+            assert((match == NULL) == (rest == NULL));
+            if (match == NULL) goto pop_2_error;
 
-                SET_TOP(rest);
-                PUSH(match);
+            if (!Py_IsNone(match)) {
                 PyErr_SetExcInfo(NULL, Py_NewRef(match), NULL);
-                Py_DECREF(exc_value);
             }
+            POKE(1, match);
+            POKE(2, rest);
             DISPATCH();
         }
 
diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
index e3d8ed340e4..3ceaca8c397 100644
--- a/Python/opcode_metadata.h
+++ b/Python/opcode_metadata.h
@@ -119,7 +119,7 @@ static const struct {
     [COMPARE_AND_BRANCH_STR] = { 2, 0, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IBC0 },
     [IS_OP] = { 2, 1, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IB },
     [CONTAINS_OP] = { 2, 1, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IB },
-    [CHECK_EG_MATCH] = { -1, -1, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IX },
+    [CHECK_EG_MATCH] = { 2, 2, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IX },
     [CHECK_EXC_MATCH] = { 2, 2, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IX },
     [IMPORT_NAME] = { 2, 1, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IB },
     [IMPORT_FROM] = { 1, 2, DIR_NONE, DIR_NONE, DIR_NONE, true, INSTR_FMT_IB },