From 1c38d896e43ca20f922f97422fb434e9a86c140a Mon Sep 17 00:00:00 2001
From: Mark Shannon <mark@hotpy.org>
Date: Fri, 3 Jan 2025 23:21:24 +0000
Subject: [PATCH] Ensure consistency between BRANCH event offsets and
 co_branches.

---
 Lib/test/test_monitoring.py | 82 +++++++++++++++++++++++++++++++++++++
 Python/bytecodes.c          |  2 -
 Python/generated_cases.c.h  |  2 -
 Python/instrumentation.c    |  6 ++-
 4 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py
index 7b9a4327a23bed..43e3e56639db62 100644
--- a/Lib/test/test_monitoring.py
+++ b/Lib/test/test_monitoring.py
@@ -1658,6 +1658,88 @@ def foo(n=0):
             exit_loop])
 
 
+class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
+
+    def check_branches(self, func, tool=TEST_TOOL, recorders=BRANCH_OFFSET_RECORDERS):
+        try:
+            self.assertEqual(sys.monitoring._all_events(), {})
+            event_list = []
+            all_events = 0
+            for recorder in recorders:
+                ev = recorder.event_type
+                sys.monitoring.register_callback(tool, ev, recorder(event_list))
+                all_events |= ev
+            sys.monitoring.set_local_events(tool, func.__code__, all_events)
+            func()
+            sys.monitoring.set_local_events(tool, func.__code__, 0)
+            for recorder in recorders:
+                sys.monitoring.register_callback(tool, recorder.event_type, None)
+            lefts = set()
+            rights = set()
+            for (src, left, right) in func.__code__.co_branches():
+                lefts.add((src, left))
+                rights.add((src, right))
+            for event in event_list:
+                way, _, src, dest = event
+                if "left" in way:
+                    self.assertIn((src, dest), lefts)
+                else:
+                    self.assertIn("right", way)
+                    self.assertIn((src, dest), rights)
+        finally:
+            sys.monitoring.set_local_events(tool, func.__code__, 0)
+            for recorder in recorders:
+                sys.monitoring.register_callback(tool, recorder.event_type, None)
+
+    def test_simple(self):
+
+        def func():
+            x = 1
+            for a in range(2):
+                if a:
+                    x = 4
+                else:
+                    x = 6
+            7
+
+        self.check_branches(func)
+
+        def whilefunc(n=0):
+            while n < 3:
+                n += 1 # line 2
+            3
+
+        self.check_branches(whilefunc)
+
+    def test_except_star(self):
+
+        class Foo:
+            def meth(self):
+                pass
+
+        def func():
+            try:
+                try:
+                    raise KeyError
+                except* Exception as e:
+                    f = Foo(); f.meth()
+            except KeyError:
+                pass
+
+
+        self.check_branches(func)
+
+    def test4(self):
+
+        def foo(n=0):
+            while n<4:
+                pass
+                n += 1
+            return None
+
+        self.check_branches(foo)
+
+
 class TestLoadSuperAttr(CheckEvents):
     RECORDERS = CallRecorder, LineRecorder, CRaiseRecorder, CReturnRecorder
 
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 3dad9383d82939..4961693c7e654a 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -362,8 +362,6 @@ dummy_func(
         }
 
         tier1 inst(INSTRUMENTED_POP_ITER, (iter -- )) {
-            // Use `this_instr+1` instead of `next_instr` as the macro assigns next_instr`.
-            (void)this_instr; // INSTRUMENTED_JUMP requires this_instr
             INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
             PyStackRef_CLOSE(iter);
         }
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 74b174e2ed87d2..b73844ca2d9542 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -4776,8 +4776,6 @@
             INSTRUCTION_STATS(INSTRUMENTED_POP_ITER);
             _PyStackRef iter;
             iter = stack_pointer[-1];
-            // Use `this_instr+1` instead of `next_instr` as the macro assigns next_instr`.
-            (void)this_instr; // INSTRUMENTED_JUMP requires this_instr
             INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
             PyStackRef_CLOSE(iter);
             stack_pointer += -1;
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index a7745377f6d23f..17e5346be5ed3d 100644
--- a/Python/instrumentation.c
+++ b/Python/instrumentation.c
@@ -3049,11 +3049,15 @@ branchesiter_next(branchesiterator *bi)
             case EXTENDED_ARG:
                 oparg = (oparg << 8) | inst.op.arg;
                 break;
+            case FOR_ITER:
+                oparg = (oparg << 8) | inst.op.arg;
+                bi->bi_offset = next_offset;
+                int target = next_offset + oparg+2; // Skips END_FOR and POP_ITER
+                return int_triple(offset*2, next_offset*2, target*2);
             case POP_JUMP_IF_FALSE:
             case POP_JUMP_IF_TRUE:
             case POP_JUMP_IF_NONE:
             case POP_JUMP_IF_NOT_NONE:
-            case FOR_ITER:
                 oparg = (oparg << 8) | inst.op.arg;
                 /* Skip NOT_TAKEN */
                 int not_taken = next_offset + 1;