GH-116017: Put JIT code and data on the same page (GH-116845)

2 files changed, 4 insertions, 5 deletions

diff --git a/Tools/jit/_stencils.py b/Tools/jit/_stencils.py
index 78c566d..05c4ce8 100644
--- a/Tools/jit/_stencils.py
+++ b/Tools/jit/_stencils.py
@@ -124,7 +124,7 @@ class Stencil:
         ):
             self.holes.append(hole.replace(offset=base + 4 * i, kind=kind))
 
-    def remove_jump(self) -> None:
+    def remove_jump(self, *, alignment: int = 1) -> None:
         """Remove a zero-length continuation jump, if it exists."""
         hole = max(self.holes, key=lambda hole: hole.offset)
         match hole:
@@ -170,7 +170,7 @@ class Stencil:
                 offset -= 2
             case _:
                 return
-        if self.body[offset:] == jump:
+        if self.body[offset:] == jump and offset % alignment == 0:
             self.body = self.body[:offset]
             self.holes.remove(hole)
 
@@ -199,9 +199,8 @@ class StencilGroup:
             ):
                 self.code.pad(alignment)
                 self.code.emit_aarch64_trampoline(hole)
-                self.code.pad(alignment)
                 self.code.holes.remove(hole)
-        self.code.remove_jump()
+        self.code.remove_jump(alignment=alignment)
         self.code.pad(alignment)
         self.data.pad(8)
         for stencil in [self.code, self.data]:
diff --git a/Tools/jit/_targets.py b/Tools/jit/_targets.py
index 417fdb5..66db358 100644
--- a/Tools/jit/_targets.py
+++ b/Tools/jit/_targets.py
@@ -89,7 +89,7 @@ class _Target(typing.Generic[_S, _R]):
         if group.data.body:
             line = f"0: {str(bytes(group.data.body)).removeprefix('b')}"
             group.data.disassembly.append(line)
-        group.process_relocations()
+        group.process_relocations(alignment=self.alignment)
         return group
 
     def _handle_section(self, section: _S, group: _stencils.StencilGroup) -> None: