Ensure consistency between BRANCH event offsets and co_branches.

Author: markshannon

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
 

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);
         }

Python/generated_cases.c.h

@@ -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;