Use NEXT_LOCATION for popping subject case

Summary: There's a new position added which is `NEXT_LOCATION` - when this is used the instruction will later be populated with the following instructions location.

Reviewed By: martindemello

Differential Revision: D80820421

fbshipit-source-id: fff6775095bec4452bb225390d3a2ab545a4aecb

Author: DinoV

cinderx/PythonLib/cinderx/compiler/pyassem.py

@@ -125,6 +125,7 @@ def __repr__(self) -> str:
 
 
 NO_LOCATION = SrcLocation(-1, -1, -1, -1)
+NEXT_LOCATION = SrcLocation(-2, -2, -2, -2)
 
 
 class Instruction:
@@ -1080,7 +1081,7 @@ def propagate_line_numbers(self) -> None:
                 continue
             prev_loc = NO_LOCATION
             for instr in block.insts:
-                if instr.lineno < 0:
+                if instr.lineno == NO_LOCATION.lineno:
                     instr.loc = prev_loc
                 else:
                     prev_loc = instr.loc
@@ -1090,7 +1091,7 @@ def propagate_line_numbers(self) -> None:
                 if next.num_predecessors == 1:
                     assert next.insts
                     next_instr = next.insts[0]
-                    if next_instr.lineno < 0:
+                    if next_instr.lineno == -1:
                         next_instr.loc = prev_loc
             last_instr = block.insts[-1]
             if (
@@ -1103,7 +1104,7 @@ def propagate_line_numbers(self) -> None:
                 if target.num_predecessors == 1:
                     assert target.insts
                     next_instr = target.insts[0]
-                    if next_instr.lineno < 0:
+                    if next_instr.lineno == NO_LOCATION.lineno:
                         next_instr.loc = prev_loc
 
     def guarantee_lineno_for_exits(self) -> None:
@@ -2061,6 +2062,16 @@ def make_line_table(self) -> bytes:
 
         loc = NO_LOCATION
         size = 0
+        prev = None
+        for t in reversed(self.insts):
+            if t.loc is NEXT_LOCATION:
+                if t.is_jump(self.opcode) or t.opname in SCOPE_EXIT_OPCODES:
+                    t.loc = NO_LOCATION
+                else:
+                    assert prev is not None
+                    t.loc = prev.loc
+            prev = t
+
         for t in self.insts:
             if t.loc != loc:
                 lpostab.emit_location(loc, size)

cinderx/PythonLib/cinderx/compiler/pycodegen.py

@@ -55,6 +55,7 @@
         FVC_STR,
         FVS_HAVE_SPEC,
         Instruction,
+        NEXT_LOCATION,
         NO_LOCATION,
         PyFlowGraph,
         PyFlowGraph310,
@@ -1827,6 +1828,10 @@ def emit_match_jump_to_end(self, end: Block) -> None:
     def emit_match_jump_to_fail_pop_unconditional(self, pc: PatternContext) -> None:
         raise NotImplementedError()
 
+    def emit_pop_subject(self, case: ast.match_case) -> None:
+        self.set_pos(case.pattern)
+        self.emit("POP_TOP")
+
     def visitMatch(self, node: ast.Match) -> None:
         """See compiler_match_inner in compile.c"""
         pc = self.pattern_context()
@@ -1863,8 +1868,7 @@ def visitMatch(self, node: ast.Match) -> None:
                 self.compileJumpIf(guard, pc.fail_pop[0], False)
             # Success! Pop the subject off, we're done with it:
             if not is_last_non_default_case:
-                self.set_pos(case.pattern)
-                self.emit("POP_TOP")
+                self.emit_pop_subject(case)
             self.visit_list(case.body)
             self.emit_match_jump_to_end(end)
             # If the pattern fails to match, we want the line number of the
@@ -6418,6 +6422,10 @@ def emit_init_class_attrs(
         gen.emit("LOAD_CONST", first_lineno)
         gen.storeName("__firstlineno__")
 
+    def emit_pop_subject(self, case: ast.match_case) -> None:
+        # Use the next location to give better locations for branch events
+        self.graph.emit_with_loc("POP_TOP", 0, NEXT_LOCATION)
+
     def maybe_add_static_attribute_to_class(self, node: ast.Attribute) -> None:
         attr_value = node.value
         if (