Better basic-block line tagging

rocky · Nov 18, 2024 · f5bc8ec · f5bc8ec
1 parent 102839e
commit f5bc8ec
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/control_flow/bb.py b/control_flow/bb.py
@@ -391,6 +391,7 @@ def basic_blocks(
     end_try_offset = None
     loop_offset = None
     return_blocks = []
+    last_line_number = code.co_firstlineno
 
     for i, inst in enumerate(instructions):
         if print_instructions:
@@ -427,7 +428,7 @@ def basic_blocks(
                 follow_offset,
                 flags,
                 jump_offsets,
-                inst.starts_line,
+                last_line_number,
             )
             loop_offset = None
             if BB_TRY in block.flags:
@@ -446,7 +447,7 @@ def basic_blocks(
                     end_offset,
                     flags,
                     jump_offsets,
-                    inst.starts_line,
+                    last_line_number,
                 )
                 loop_offset = None
                 if BB_TRY in block.flags:
@@ -458,6 +459,11 @@ def basic_blocks(
             if offset in loop_targets:
                 flags.add(BB_LOOP)
 
+        # This should be done after closing off the
+        # basic block above.
+        if inst.starts_line is not None:
+            last_line_number = inst.starts_line
+
         # Add block flags for certain classes of instructions
         if op in bb.JUMP_CONDITIONAL:
             flags.add(BB_JUMP_CONDITIONAL)
@@ -499,7 +505,7 @@ def basic_blocks(
                 follow_offset,
                 flags,
                 jump_offsets,
-                inst.starts_line,
+                last_line_number,
             )
             loop_offset = None
             start_offset = follow_offset
@@ -529,7 +535,7 @@ def basic_blocks(
                     follow_offset,
                     flags,
                     jump_offsets,
-                    inst.starts_line,
+                    last_line_number,
                 )
                 loop_offset = None
                 if BB_TRY in block.flags:
@@ -548,7 +554,7 @@ def basic_blocks(
                     follow_offset,
                     flags,
                     jump_offsets,
-                    inst.starts_line,
+                    last_line_number,
                 )
                 loop_offset = None
                 if BB_TRY in block.flags:
@@ -567,7 +573,7 @@ def basic_blocks(
                 follow_offset,
                 flags,
                 jump_offsets,
-                inst.starts_line,
+                last_line_number,
             )
             loop_offset = None
             start_offset = follow_offset