From e4d3a96a113070fde433834a6c9fb79ebeebad4a Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@microsoft.com>
Date: Fri, 22 Jul 2022 16:28:03 -0700
Subject: [PATCH] GH-94438: Handle extended arguments and conditional pops in
 mark_stacks (GH-95110)

---
 Lib/test/test_sys_settrace.py                 | 36 +++++++++++++++++++
 ...2-07-22-12-53-34.gh-issue-94438.hNqACc.rst |  4 +++
 Objects/frameobject.c                         | 13 ++++---
 3 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2022-07-22-12-53-34.gh-issue-94438.hNqACc.rst

diff --git a/Lib/test/test_sys_settrace.py b/Lib/test/test_sys_settrace.py
index 8d0c3171441..9f1aa81dbcd 100644
--- a/Lib/test/test_sys_settrace.py
+++ b/Lib/test/test_sys_settrace.py
@@ -2685,6 +2685,42 @@ def test_jump_with_null_on_stack_load_attr(output):
         )
         output.append(15)
 
+    @jump_test(2, 3, [1, 3])
+    def test_jump_extended_args_unpack_ex_simple(output):
+        output.append(1)
+        _, *_, _ = output.append(2) or "Spam"
+        output.append(3)
+
+    @jump_test(3, 4, [1, 4, 4, 5])
+    def test_jump_extended_args_unpack_ex_tricky(output):
+        output.append(1)
+        (
+            _, *_, _
+        ) = output.append(4) or "Spam"
+        output.append(5)
+
+    def test_jump_extended_args_for_iter(self):
+        # In addition to failing when extended arg handling is broken, this can
+        # also hang for a *very* long time:
+        source = [
+            "def f(output):",
+            "    output.append(1)",
+            "    for _ in spam:",
+            *(f"        output.append({i})" for i in range(3, 100_000)),
+            f"    output.append(100_000)",
+        ]
+        namespace = {}
+        exec("\n".join(source), namespace)
+        f = namespace["f"]
+        self.run_test(f,  2, 100_000, [1, 100_000])
+
+    @jump_test(2, 3, [1, 3])
+    def test_jump_or_pop(output):
+        output.append(1)
+        _ = output.append(2) and "Spam"
+        output.append(3)
+
+
 class TestExtendedArgs(unittest.TestCase):
 
     def setUp(self):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-07-22-12-53-34.gh-issue-94438.hNqACc.rst b/Misc/NEWS.d/next/Core and Builtins/2022-07-22-12-53-34.gh-issue-94438.hNqACc.rst
new file mode 100644
index 00000000000..2a7249a833c
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-07-22-12-53-34.gh-issue-94438.hNqACc.rst	
@@ -0,0 +1,4 @@
+Fix an issue that caused extended opcode arguments and some conditional pops
+to be ignored when calculating valid jump targets for assignments to the
+``f_lineno`` attribute of frame objects. In some cases, this could cause
+inconsistent internal state, resulting in a hard crash of the interpreter.
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 9e5450af52a..26b38bae780 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -319,11 +319,15 @@ mark_stacks(PyCodeObject *code_obj, int len)
                     int64_t target_stack;
                     int j = get_arg(code, i);
                     if (opcode == POP_JUMP_FORWARD_IF_FALSE ||
-                        opcode == POP_JUMP_FORWARD_IF_TRUE) {
+                        opcode == POP_JUMP_FORWARD_IF_TRUE ||
+                        opcode == JUMP_IF_FALSE_OR_POP ||
+                        opcode == JUMP_IF_TRUE_OR_POP)
+                    {
                         j += i + 1;
                     }
-                    else if (opcode == POP_JUMP_BACKWARD_IF_FALSE ||
-                             opcode == POP_JUMP_BACKWARD_IF_TRUE) {
+                    else {
+                        assert(opcode == POP_JUMP_BACKWARD_IF_FALSE ||
+                               opcode == POP_JUMP_BACKWARD_IF_TRUE);
                         j = i + 1 - j;
                     }
                     assert(j < len);
@@ -459,7 +463,8 @@ mark_stacks(PyCodeObject *code_obj, int len)
                 }
                 default:
                 {
-                    int delta = PyCompile_OpcodeStackEffect(opcode, _Py_OPARG(code[i]));
+                    int delta = PyCompile_OpcodeStackEffect(opcode, get_arg(code, i));
+                    assert(delta != PY_INVALID_STACK_EFFECT);
                     while (delta < 0) {
                         next_stack = pop_value(next_stack);
                         delta++;
-- 
GitLab