From 7ab3ec65ed4947af23b13df8827fef2e9aecee93 Mon Sep 17 00:00:00 2001
From: Matt Page <mpage@meta.com>
Date: Fri, 10 Jan 2025 12:42:47 -0800
Subject: [PATCH] Make POP_INPUT take the name of the TOS

---
 Lib/test/test_generated_cases.py           | 15 +++++++++++++--
 Python/bytecodes.c                         | 12 ++++++------
 Tools/cases_generator/generators_common.py |  7 ++++++-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/Lib/test/test_generated_cases.py b/Lib/test/test_generated_cases.py
index ae1dc90b18efef..ced4e0c99f3d7f 100644
--- a/Lib/test/test_generated_cases.py
+++ b/Lib/test/test_generated_cases.py
@@ -1677,7 +1677,7 @@ def test_pystackref_frompyobject_new_next_to_cmacro(self):
     def test_pop_input(self):
         input = """
         inst(OP, (a, b --)) {
-            POP_INPUT();
+            POP_INPUT(b);
             HAM(a);
             INPUTS_DEAD();
         }
@@ -1688,6 +1688,8 @@ def test_pop_input(self):
             next_instr += 1;
             INSTRUCTION_STATS(OP);
             _PyStackRef a;
+            _PyStackRef b;
+            b = stack_pointer[-1];
             a = stack_pointer[-2];
             stack_pointer += -1;
             assert(WITHIN_STACK_BOUNDS());
@@ -1702,7 +1704,16 @@ def test_pop_input(self):
     def test_pop_input_with_empty_stack(self):
         input = """
         inst(OP, (--)) {
-            POP_INPUT();
+            POP_INPUT(foo);
+        }
+        """
+        with self.assertRaises(SyntaxError):
+            self.run_cases_test(input, "")
+
+    def test_pop_input_with_non_tos(self):
+        input = """
+        inst(OP, (a, b --)) {
+            POP_INPUT(a);
         }
         """
         with self.assertRaises(SyntaxError):
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 79b659131fcd9b..e7b1255cf5233c 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -2222,7 +2222,7 @@ dummy_func(
             PyDictUnicodeEntry *ep = DK_UNICODE_ENTRIES(mod_keys) + index;
             PyObject *attr_o = FT_ATOMIC_LOAD_PTR_RELAXED(ep->me_value);
             // Clear mod_keys from stack in case we need to deopt
-            POP_INPUT();
+            POP_INPUT(mod_keys);
             DEOPT_IF(attr_o == NULL);
             #ifdef Py_GIL_DISABLED
             int increfed = _Py_TryIncrefCompareStackRef(&ep->me_value, attr_o, &attr);
@@ -2257,31 +2257,31 @@ dummy_func(
         op(_LOAD_ATTR_WITH_HINT, (hint/1, owner, dict: PyDictObject * -- attr, null if (oparg & 1))) {
             PyObject *attr_o;
             if (!LOCK_OBJECT(dict)) {
-                POP_INPUT();
+                POP_INPUT(dict);
                 DEOPT_IF(true);
             }
 
             if (hint >= (size_t)dict->ma_keys->dk_nentries) {
                 UNLOCK_OBJECT(dict);
-                POP_INPUT();
+                POP_INPUT(dict);
                 DEOPT_IF(true);
             }
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg>>1);
             if (dict->ma_keys->dk_kind != DICT_KEYS_UNICODE) {
                 UNLOCK_OBJECT(dict);
-                POP_INPUT();
+                POP_INPUT(dict);
                 DEOPT_IF(true);
             }
             PyDictUnicodeEntry *ep = DK_UNICODE_ENTRIES(dict->ma_keys) + hint;
             if (ep->me_key != name) {
                 UNLOCK_OBJECT(dict);
-                POP_INPUT();
+                POP_INPUT(dict);
                 DEOPT_IF(true);
             }
             attr_o = ep->me_value;
             if (attr_o == NULL) {
                 UNLOCK_OBJECT(dict);
-                POP_INPUT();
+                POP_INPUT(dict);
                 DEOPT_IF(true);
             }
             STAT_INC(LOAD_ATTR, hit);
diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py
index 59199c870eacf4..01127f81e34cec 100644
--- a/Tools/cases_generator/generators_common.py
+++ b/Tools/cases_generator/generators_common.py
@@ -358,11 +358,16 @@ def pop_input(
         inst: Instruction | None,
     ) -> bool:
         next(tkn_iter)
+        name_tkn = next(tkn_iter)
+        name = name_tkn.text
         next(tkn_iter)
         next(tkn_iter)
         if not storage.inputs:
             raise analysis_error("stack is empty", tkn)
-        storage.inputs[-1].defined = False
+        tos = storage.inputs[-1]
+        if tos.name != name:
+            raise analysis_error(f"'{name} is not top of stack", name_tkn)
+        tos.defined = False
         storage.clear_dead_inputs()
         storage.flush(self.out)
         return True