[Python-checkins] GH-103082: Turn on branch events for FOR_ITER instructions. (#103507)

[markshannon]

https://github.com/python/cpython/commit/70e0a28bed1354ec450989271152ae697e1f0993
commit: 70e0a28bed1354ec450989271152ae697e1f0993
branch: main
author: Mark Shannon <mark at hotpy.org>
committer: markshannon <mark at hotpy.org>
date: 2023-04-13T13:56:09+01:00
summary:

GH-103082: Turn on branch events for FOR_ITER instructions. (#103507)

Turn on branch events for FOR_ITER instructions.

files:
M Lib/test/test_monitoring.py
M Python/instrumentation.c

diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py
index 4aad3da61b6b..738ace923cc5 100644
--- a/Lib/test/test_monitoring.py
+++ b/Lib/test/test_monitoring.py
@@ -1011,6 +1011,77 @@ def func3():
             ('line', 'func3', 6)])
 
 
+def line_from_offset(code, offset):
+    for start, end, line in code.co_lines():
+        if start <= offset < end:
+            return line - code.co_firstlineno
+    return -1
+
+class JumpRecorder:
+
+    event_type = E.JUMP
+    name = "jump"
+
+    def __init__(self, events):
+        self.events = events
+
+    def __call__(self, code, from_, to):
+        from_line = line_from_offset(code, from_)
+        to_line = line_from_offset(code, to)
+        self.events.append((self.name, code.co_name, from_line, to_line))
+
+
+class BranchRecorder(JumpRecorder):
+
+    event_type = E.BRANCH
+    name = "branch"
+
+
+JUMP_AND_BRANCH_RECORDERS = JumpRecorder, BranchRecorder
+JUMP_BRANCH_AND_LINE_RECORDERS = JumpRecorder, BranchRecorder, LineRecorder
+
+class TestBranchAndJumpEvents(CheckEvents):
+    maxDiff = None
+
+    def test_loop(self):
+
+        def func():
+            x = 1
+            for a in range(2):
+                if a:
+                    x = 4
+                else:
+                    x = 6
+
+        self.check_events(func, recorders = JUMP_AND_BRANCH_RECORDERS, expected = [
+            ('branch', 'func', 2, 2),
+            ('branch', 'func', 3, 6),
+            ('jump', 'func', 6, 2),
+            ('branch', 'func', 2, 2),
+            ('branch', 'func', 3, 4),
+            ('jump', 'func', 4, 2),
+            ('branch', 'func', 2, 2)])
+
+
+        self.check_events(func, recorders = JUMP_BRANCH_AND_LINE_RECORDERS, expected = [
+            ('line', 'check_events', 10),
+            ('line', 'func', 1),
+            ('line', 'func', 2),
+            ('branch', 'func', 2, 2),
+            ('line', 'func', 3),
+            ('branch', 'func', 3, 6),
+            ('line', 'func', 6),
+            ('jump', 'func', 6, 2),
+            ('branch', 'func', 2, 2),
+            ('line', 'func', 3),
+            ('branch', 'func', 3, 4),
+            ('line', 'func', 4),
+            ('jump', 'func', 4, 2),
+            ('branch', 'func', 2, 2),
+            ('line', 'func', 2),
+            ('line', 'check_events', 11)])
+
+
 class TestSetGetEvents(MonitoringTestBase, unittest.TestCase):
 
     def test_global(self):
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index 39a7eaa3bd2d..8dc8b01fcb04 100644
--- a/Python/instrumentation.c
+++ b/Python/instrumentation.c
@@ -105,6 +105,8 @@ static const uint8_t INSTRUMENTED_OPCODES[256] = {
     [INSTRUMENTED_END_FOR] = INSTRUMENTED_END_FOR,
     [END_SEND] = INSTRUMENTED_END_SEND,
     [INSTRUMENTED_END_SEND] = INSTRUMENTED_END_SEND,
+    [FOR_ITER] = INSTRUMENTED_FOR_ITER,
+    [INSTRUMENTED_FOR_ITER] = INSTRUMENTED_FOR_ITER,
 
     [INSTRUMENTED_LINE] = INSTRUMENTED_LINE,
     [INSTRUMENTED_INSTRUCTION] = INSTRUMENTED_INSTRUCTION,