a proper fix by special casing calls

Author: Fidget-Spinner

Tools/jit/_optimizers.py

@@ -95,6 +95,7 @@ class InstructionKind(enum.Enum):
     JUMP = enum.auto()
     LONG_BRANCH = enum.auto()
     SHORT_BRANCH = enum.auto()
+    CALL = enum.auto()
     RETURN = enum.auto()
     SMALL_CONST_1 = enum.auto()
     SMALL_CONST_2 = enum.auto()
@@ -182,6 +183,8 @@ class Optimizer:
     # Two groups (instruction and target):
     _re_branch: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH
     # One group (target):
+    _re_call: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH
+    # One group (target):
     _re_jump: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH
     # No groups:
     _re_return: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH
@@ -220,7 +223,7 @@ def __post_init__(self) -> None:
                 block.link = block = _Block()
             inst = self._parse_instruction(line)
             block.instructions.append(inst)
-            if inst.is_branch():
+            if inst.is_branch() or inst.kind == InstructionKind.CALL:
                 # A block ending in a branch has a target and fallthrough:
                 assert inst.target is not None
                 block.target = self._lookup_label(inst.target)
@@ -256,6 +259,10 @@ def _parse_instruction(self, line: str) -> Instruction:
             target = match["target"]
             name = line[: -len(target)].strip()
             kind = InstructionKind.JUMP
+        elif match := self._re_call.match(line):
+            target = match["target"]
+            name = line[: -len(target)].strip()
+            kind = InstructionKind.JUMP
         elif match := self._re_return.match(line):
             name = line
             kind = InstructionKind.RETURN
@@ -566,6 +573,8 @@ class OptimizerAArch64(Optimizer):  # pylint: disable = too-few-public-methods
         rf"\s*(?P<instruction>{'|'.join(_branch_patterns)})\s+(.+,\s+)*(?P<target>[\w.]+)"
     )
 
+    # https://developer.arm.com/documentation/ddi0406/b/Application-Level-Architecture/Instruction-Details/Alphabetical-list-of-instructions/BL--BLX--immediate-
+    _re_call = re.compile(r"\s*blx??\s+(?P<target>[\w.]+)")
     # https://developer.arm.com/documentation/ddi0602/2025-03/Base-Instructions/B--Branch-
     _re_jump = re.compile(r"\s*b\s+(?P<target>[\w.]+)")
     # https://developer.arm.com/documentation/ddi0602/2025-09/Base-Instructions/RET--Return-from-subroutine-
@@ -628,9 +637,9 @@ class OptimizerX86(Optimizer):  # pylint: disable = too-few-public-methods
     _re_branch = re.compile(
         rf"\s*(?P<instruction>{'|'.join(_X86_BRANCHES)})\s+(?P<target>[\w.]+)"
     )
-    # https://www.felixcloutier.com/x86/jmp
     # https://www.felixcloutier.com/x86/call
-    # Calls are also logically jumps to labels.
-    _re_jump = re.compile(r"\s*((?:jmp)|(?:callq?))\s+(?P<target>[\w.]+)")
+    _re_call = re.compile(r"\s*callq?\s+(?P<target>[\w.]+)")
+    # https://www.felixcloutier.com/x86/jmp
+    _re_jump = re.compile(r"\s*jmp\s+(?P<target>[\w.]+)")
     # https://www.felixcloutier.com/x86/ret
     _re_return = re.compile(r"\s*ret\b")