Trampoline attempt

Author: savannahostrowski

Python/jit.c

@@ -216,11 +216,7 @@ patch_32r(unsigned char *location, uint64_t value)
     value -= (uintptr_t)location;
     // Check that we're not out of range of 32 signed bits:
     assert((int64_t)value >= -(1LL << 31));
-    // assert((int64_t)value < (1LL << 31));
-    if ((int64_t)value >= (1LL << 31)) {
-        __builtin_debugtrap();
-    }
-
+    assert((int64_t)value < (1LL << 31));
     *loc32 = (uint32_t)value;
 
 }
@@ -424,12 +420,17 @@ patch_x86_64_32rx(unsigned char *location, uint64_t value)
 }
 
 void patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *state);
+void patch_x86_64_trampoline(unsigned char *location, int ordinal, jit_state *state);
 
 #include "jit_stencils.h"
 
 #if defined(__aarch64__) || defined(_M_ARM64)
     #define TRAMPOLINE_SIZE 16
     #define DATA_ALIGN 8
+#elif defined(__x86_64__) && defined(__APPLE__)
+    // x86_64 trampolines: jmp *(%rip); .quad address (6 bytes + 8 bytes = 14 bytes)
+    #define TRAMPOLINE_SIZE 16  // Round up to 16 for alignment
+    #define DATA_ALIGN 16
 #else
     #define TRAMPOLINE_SIZE 0
     #define DATA_ALIGN 1
@@ -481,6 +482,47 @@ patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *state)
     patch_aarch64_26r(location, (uintptr_t)p);
 }
 
+// Generate and patch x86_64 trampolines.
+void
+patch_x86_64_trampoline(unsigned char *location, int ordinal, jit_state *state)
+{
+    uint64_t value = (uintptr_t)symbols_map[ordinal];
+    int64_t range = (int64_t)value - 4 - (int64_t)location;
+
+    // If we are in range of 32 signed bits, patch directly
+    if (range >= -(1LL << 31) && range < (1LL << 31)) {
+        patch_32r(location, value - 4);
+        return;
+    }
+
+    // Out of range - need a trampoline
+    const uint32_t symbol_mask = 1 << (ordinal % 32);
+    const uint32_t trampoline_mask = state->trampolines.mask[ordinal / 32];
+    assert(symbol_mask & trampoline_mask);
+
+    // Count the number of set bits in the trampoline mask lower than ordinal
+    int index = _Py_popcount32(trampoline_mask & (symbol_mask - 1));
+    for (int i = 0; i < ordinal / 32; i++) {
+        index += _Py_popcount32(state->trampolines.mask[i]);
+    }
+
+    unsigned char *trampoline = state->trampolines.mem + index * TRAMPOLINE_SIZE;
+    assert((size_t)(index + 1) * TRAMPOLINE_SIZE <= state->trampolines.size);
+
+    /* Generate the trampoline (14 bytes, padded to 16):
+       0: ff 25 00 00 00 00    jmp *(%rip)  # Jump to address at offset 6
+       6: XX XX XX XX XX XX XX XX   .quad value (64-bit address)
+    */
+    trampoline[0] = 0xFF;  // jmp opcode
+    trampoline[1] = 0x25;  // ModRM byte for jmp *disp32(%rip)
+    // Offset 0: the address is right after this instruction (at offset 6)
+    *(uint32_t *)(trampoline + 2) = 0;
+    *(uint64_t *)(trampoline + 6) = value;
+
+    // Patch the call site to call the trampoline instead
+    patch_32r(location, (uintptr_t)trampoline - 4);
+}
+
 static void
 combine_symbol_mask(const symbol_mask src, symbol_mask dest)
 {

Tools/jit/_stencils.py

@@ -239,6 +239,23 @@ def process_relocations(self, known_symbols: dict[str, int]) -> None:
                 self._trampolines.add(ordinal)
                 hole.addend = ordinal
                 hole.symbol = None
+            # x86_64 Darwin trampolines for external symbols
+            elif (
+                hole.kind == "X86_64_RELOC_BRANCH"
+                and hole.value is HoleValue.ZERO
+                and hole.symbol not in self.symbols
+            ):
+                hole.func = "patch_x86_64_trampoline"
+                hole.need_state = True
+                assert hole.symbol is not None
+                if hole.symbol in known_symbols:
+                    ordinal = known_symbols[hole.symbol]
+                else:
+                    ordinal = len(known_symbols)
+                    known_symbols[hole.symbol] = ordinal
+                self._trampolines.add(ordinal)
+                hole.addend = ordinal
+                hole.symbol = None
         self.data.pad(8)
         for stencil in [self.code, self.data]:
             for hole in stencil.holes: