add fix suggested by Claude

Author: kumaraditya303

Python/jit.c

@@ -554,6 +554,7 @@ patch_x86_64_32rx(unsigned char *location, uint64_t value)
 
 void patch_got_symbol(jit_state *state, int ordinal);
 void patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *state);
+void patch_aarch64_trampoline_addr(unsigned char *location, int ordinal, uint64_t value, jit_state *state);
 void patch_x86_64_trampoline(unsigned char *location, int ordinal, jit_state *state);
 
 #include "jit_stencils.h"
@@ -585,28 +586,27 @@ patch_got_symbol(jit_state *state, int ordinal)
 void
 patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *state)
 {
-
     uint64_t value = (uintptr_t)symbols_map[ordinal];
-    int64_t range = value - (uintptr_t)location;
+    patch_aarch64_trampoline_addr(location, ordinal, value, state);
+}
+
+// Generate and patch AArch64 trampolines for dynamic addresses (e.g. operands).
+// Unlike patch_aarch64_trampoline, the target address is passed directly rather
+// than looked up from symbols_map. The ordinal is used to allocate a trampoline slot.
+void
+patch_aarch64_trampoline_addr(unsigned char *location, int ordinal, uint64_t value, jit_state *state)
+{
+    int64_t range = (int64_t)value - (int64_t)(uintptr_t)location;
 
-    // If we are in range of 28 signed bits, we patch the instruction with
-    // the address of the symbol.
     if (range >= -(1 << 27) && range < (1 << 27)) {
-        patch_aarch64_26r(location, (uintptr_t)value);
+        patch_aarch64_26r(location, value);
         return;
     }
 
-    // Out of range - need a trampoline
     uint32_t *p = (uint32_t *)get_symbol_slot(ordinal, &state->trampolines, TRAMPOLINE_SIZE);
 
-    /* Generate the trampoline
-       0: 58000048      ldr     x8, 8
-       4: d61f0100      br      x8
-       8: 00000000      // The next two words contain the 64-bit address to jump to.
-       c: 00000000
-    */
-    p[0] = 0x58000048;
-    p[1] = 0xD61F0100;
+    p[0] = 0x58000048; // ldr x8, 8
+    p[1] = 0xD61F0100; // br x8
     p[2] = value & 0xffffffff;
     p[3] = value >> 32;
 

Tools/jit/_stencils.py

@@ -296,6 +296,23 @@ def process_relocations(self, known_symbols: dict[str, int]) -> None:
                 self._trampolines.add(ordinal)
                 hole.addend = ordinal
                 hole.symbol = None
+            # aarch64 trampolines for operand-based call targets (e.g. inlined cfunc)
+            elif (
+                hole.kind
+                in {"R_AARCH64_CALL26", "R_AARCH64_JUMP26", "ARM64_RELOC_BRANCH26"}
+                and hole.value in {HoleValue.OPERAND0, HoleValue.OPERAND1}
+            ):
+                value_expr = _HOLE_EXPRS[hole.value]
+                synth_name = f"_JIT_TRAMPOLINE_{hole.value.name}"
+                if synth_name in known_symbols:
+                    ordinal = known_symbols[synth_name]
+                else:
+                    ordinal = len(known_symbols)
+                    known_symbols[synth_name] = ordinal
+                self._trampolines.add(ordinal)
+                hole.func = "patch_aarch64_trampoline_addr"
+                hole.need_state = True
+                hole.custom_value = f"{ordinal}, {value_expr}"
             # x86_64 Darwin trampolines for external symbols
             elif (
                 hole.kind == "X86_64_RELOC_BRANCH"

Tools/jit/_writer.py

@@ -35,7 +35,10 @@ def _dump_footer(
     yield f"static const void * const symbols_map[{max(len(symbols), 1)}] = {{"
     if symbols:
         for symbol, ordinal in symbols.items():
-            yield f"    [{ordinal}] = &{symbol},"
+            if symbol.startswith("_JIT_TRAMPOLINE_"):
+                yield f"    [{ordinal}] = 0,"
+            else:
+                yield f"    [{ordinal}] = &{symbol},"
     else:
         yield "    0"
     yield "};"