gh-128627: Use __builtin_wasm_test_function_pointer_signature for Emscripten trampolines

hoodmane · hoodmane · commit 6bbee1e51bfd · 2025-08-06T17:28:58.000+02:00
Since llvm/llvm-project#150201 was merged, there is now a better way to do this. Requires Emscripten 4.0.12.
diff --git a/Include/internal/pycore_runtime_structs.h b/Include/internal/pycore_runtime_structs.h
@@ -279,7 +279,10 @@ struct pyruntimestate {
 #if defined(__EMSCRIPTEN__) && defined(PY_CALL_TRAMPOLINE)
     // Used in "Python/emscripten_trampoline.c" to choose between type
     // reflection trampoline and EM_JS trampoline.
-    int (*emscripten_count_args_function)(PyCFunctionWithKeywords func);
+    PyObject *(*emscripten_trampoline)(int *success,
+                                       PyCFunctionWithKeywords func,
+                                       PyObject *arg1, PyObject *arg2,
+                                       PyObject *arg3);
 #endif
 
     /* All the objects that are shared by the runtime's interpreters. */
diff --git a/Makefile.pre.in b/Makefile.pre.in
@@ -3106,6 +3106,11 @@ config.status:	$(srcdir)/configure
 Python/asm_trampoline.o: $(srcdir)/Python/asm_trampoline.S
 	$(CC) -c $(PY_CORE_CFLAGS) -o $@ $<
 
+Python/emscripten_trampoline_inner.wasm: $(srcdir)/Python/emscripten_trampoline_inner.c
+	$$(dirname $$(dirname $(CC)))/bin/clang -o $@ $< -mgc -O2 -Wl,--no-entry -Wl,--import-table -Wl,--import-memory -target wasm32-unknown-unknown -nostdlib
+
+Python/emscripten_trampoline.o: $(srcdir)/Python/emscripten_trampoline.c Python/emscripten_trampoline_inner.wasm
+	$(CC) -c $(PY_CORE_CFLAGS) -o $@ $<
 
 JIT_DEPS = \
 		$(srcdir)/Tools/jit/*.c \
diff --git a/Python/emscripten_trampoline.c b/Python/emscripten_trampoline.c
@@ -4,7 +4,16 @@
 #include <Python.h>
 #include "pycore_runtime.h"         // _PyRuntime
 
-typedef int (*CountArgsFunc)(PyCFunctionWithKeywords func);
+typedef PyObject *(*TrampolineFunc)(int *success, PyCFunctionWithKeywords func,
+                                    PyObject *arg1, PyObject *arg2,
+                                    PyObject *arg3);
+
+EMSCRIPTEN_KEEPALIVE const char trampoline_inner_wasm[] = {
+#embed "Python/emscripten_trampoline_inner.wasm"
+};
+
+EMSCRIPTEN_KEEPALIVE int trampoline_inner_wasm_length =
+    sizeof(trampoline_inner_wasm);
 
 // Offset of emscripten_count_args_function in _PyRuntimeState. There's a couple
 // of alternatives:
@@ -18,159 +27,55 @@ typedef int (*CountArgsFunc)(PyCFunctionWithKeywords func);
 //
 // So putting the mutable constant in _PyRuntime and using a immutable global to
 // record the offset so we can access it from JS is probably the best way.
-EMSCRIPTEN_KEEPALIVE const int _PyEM_EMSCRIPTEN_COUNT_ARGS_OFFSET = offsetof(_PyRuntimeState, emscripten_count_args_function);
-
-EM_JS(CountArgsFunc, _PyEM_GetCountArgsPtr, (), {
-    return Module._PyEM_CountArgsPtr; // initialized below
-}
-// Binary module for the checks. It has to be done in web assembly because
-// clang/llvm have no support yet for the reference types yet. In fact, the wasm
-// binary toolkit doesn't yet support the ref.test instruction either. To
-// convert the following textual wasm to a binary, you can build wabt from this
-// branch: https://github.com/WebAssembly/wabt/pull/2529 and then use that
-// wat2wasm binary.
-//
-// (module
-//     (type $type0 (func (param) (result i32)))
-//     (type $type1 (func (param i32) (result i32)))
-//     (type $type2 (func (param i32 i32) (result i32)))
-//     (type $type3 (func (param i32 i32 i32) (result i32)))
-//     (type $blocktype (func (param) (result)))
-//     (table $funcs (import "e" "t") 0 funcref)
-//     (export "f" (func $f))
-//     (func $f (param $fptr i32) (result i32)
-//              (local $fref funcref)
-//         local.get $fptr
-//         table.get $funcs
-//         local.tee $fref
-//         ref.test $type3
-//         if $blocktype
-//             i32.const 3
-//             return
-//         end
-//         local.get $fref
-//         ref.test $type2
-//         if $blocktype
-//             i32.const 2
-//             return
-//         end
-//         local.get $fref
-//         ref.test $type1
-//         if $blocktype
-//             i32.const 1
-//             return
-//         end
-//         local.get $fref
-//         ref.test $type0
-//         if $blocktype
-//             i32.const 0
-//             return
-//         end
-//         i32.const -1
-//     )
-// )
-
-function getPyEMCountArgsPtr() {
-    // Starting with iOS 18.3.1, WebKit on iOS has an issue with the garbage
-    // collector that breaks the call trampoline. See #130418 and
-    // https://bugs.webkit.org/show_bug.cgi?id=293113 for details.
-    let isIOS = globalThis.navigator && (
-        /iPad|iPhone|iPod/.test(navigator.userAgent) ||
-        // Starting with iPadOS 13, iPads might send a platform string that looks like a desktop Mac.
-        // To differentiate, we check if the platform is 'MacIntel' (common for Macs and newer iPads)
-        // AND if the device has multi-touch capabilities (navigator.maxTouchPoints > 1)
-        (navigator.platform === 'MacIntel' && typeof navigator.maxTouchPoints !== 'undefined' && navigator.maxTouchPoints > 1)
-    );
-    if (isIOS) {
+EMSCRIPTEN_KEEPALIVE const int _PyEM_EMSCRIPTEN_TRAMPOLINE_OFFSET =
+    offsetof(_PyRuntimeState, emscripten_trampoline);
+
+EM_JS(
+    TrampolineFunc, _PyEM_GetTrampolinePtr, (void),
+    {
+      return Module._PyEM_CountArgsPtr; // initialized below
+    } function getPyEMTrampolinePtr() {
+      // Starting with iOS 18.3.1, WebKit on iOS has an issue with the garbage
+      // collector that breaks the call trampoline. See #130418 and
+      // https://bugs.webkit.org/show_bug.cgi?id=293113 for details.
+      let isIOS = globalThis.navigator &&
+                  (/ iPad | iPhone | iPod /.test(navigator.userAgent) ||
+                   // Starting with iPadOS 13, iPads might send a platform
+                   // string that looks like a desktop Mac. To differentiate, we
+                   // check if the platform is 'MacIntel' (common for Macs and
+                   // newer iPads) AND if the device has multi-touch
+                   // capabilities (navigator.maxTouchPoints > 1)
+                   (navigator.platform ==
+                    = 'MacIntel' &&typeof navigator.maxTouchPoints !=
+                    = 'undefined' && navigator.maxTouchPoints > 1));
+      if (isIOS) {
         return 0;
-    }
-
-    // Try to initialize countArgsFunc
-    const code = new Uint8Array([
-        0x00, 0x61, 0x73, 0x6d, // \0asm magic number
-        0x01, 0x00, 0x00, 0x00, // version 1
-        0x01, 0x1a, // Type section, body is 0x1a bytes
-            0x05, // 6 entries
-            0x60, 0x00, 0x01, 0x7f,                      // (type $type0 (func (param) (result i32)))
-            0x60, 0x01, 0x7f, 0x01, 0x7f,                // (type $type1 (func (param i32) (result i32)))
-            0x60, 0x02, 0x7f, 0x7f, 0x01, 0x7f,          // (type $type2 (func (param i32 i32) (result i32)))
-            0x60, 0x03, 0x7f, 0x7f, 0x7f, 0x01, 0x7f,    // (type $type3 (func (param i32 i32 i32) (result i32)))
-            0x60, 0x00, 0x00,                            // (type $blocktype (func (param) (result)))
-        0x02, 0x09, // Import section, 0x9 byte body
-            0x01, // 1 import (table $funcs (import "e" "t") 0 funcref)
-            0x01, 0x65, // "e"
-            0x01, 0x74, // "t"
-            0x01,       // importing a table
-            0x70,       // of entry type funcref
-            0x00, 0x00, // table limits: no max, min of 0
-        0x03, 0x02,   // Function section
-            0x01, 0x01, // We're going to define one function of type 1 (func (param i32) (result i32))
-        0x07, 0x05, // export section
-            0x01, // 1 export
-            0x01, 0x66, // called "f"
-            0x00, // a function
-            0x00, // at index 0
-
-        0x0a, 56,  // Code section,
-            0x01, 54, // one entry of length 54
-            0x01, 0x01, 0x70, // one local of type funcref
-            // Body of the function
-            0x20, 0x00,       // local.get $fptr
-            0x25, 0x00,       // table.get $funcs
-            0x22, 0x01,       // local.tee $fref
-            0xfb, 0x14, 0x03, // ref.test $type3
-            0x04, 0x04,       // if (type $blocktype)
-                0x41, 0x03,   //   i32.const 3
-                0x0f,         //   return
-            0x0b,             // end block
-
-            0x20, 0x01,       // local.get $fref
-            0xfb, 0x14, 0x02, // ref.test $type2
-            0x04, 0x04,       // if (type $blocktype)
-                0x41, 0x02,   //   i32.const 2
-                0x0f,         //   return
-            0x0b,             // end block
-
-            0x20, 0x01,       // local.get $fref
-            0xfb, 0x14, 0x01, // ref.test $type1
-            0x04, 0x04,       // if (type $blocktype)
-                0x41, 0x01,   //   i32.const 1
-                0x0f,         //   return
-            0x0b,             // end block
-
-            0x20, 0x01,       // local.get $fref
-            0xfb, 0x14, 0x00, // ref.test $type0
-            0x04, 0x04,       // if (type $blocktype)
-                0x41, 0x00,   //   i32.const 0
-                0x0f,         //   return
-            0x0b,             // end block
-
-            0x41, 0x7f,       // i32.const -1
-            0x0b // end function
-    ]);
-    try {
+      }
+      const code = HEAP8.subarray(
+          _trampoline_inner_wasm,
+          _trampoline_inner_wasm + HEAP32[_trampoline_inner_wasm_length / 4]);
+      try {
         const mod = new WebAssembly.Module(code);
         const inst = new WebAssembly.Instance(mod, { e: { t: wasmTable } });
-        return addFunction(inst.exports.f);
-    } catch (e) {
+        return addFunction(inst.exports.trampoline_call);
+      } catch (e) {
         // If something goes wrong, we'll null out _PyEM_CountFuncParams and fall
         // back to the JS trampoline.
         return 0;
+      }
     }
-}
 
-addOnPreRun(() => {
-    const ptr = getPyEMCountArgsPtr();
-    Module._PyEM_CountArgsPtr = ptr;
-    const offset = HEAP32[__PyEM_EMSCRIPTEN_COUNT_ARGS_OFFSET / 4];
-    HEAP32[(__PyRuntime + offset) / 4] = ptr;
-});
-);
+    addOnPreRun(() = > {
+      const ptr = getPyEMTrampolinePtr();
+      Module._PyEM_CountArgsPtr = ptr;
+      const offset = HEAP32[__PyEM_EMSCRIPTEN_TRAMPOLINE_OFFSET / 4];
+      HEAP32[(__PyRuntime + offset) / 4] = ptr;
+    }););
 
 void
 _Py_EmscriptenTrampoline_Init(_PyRuntimeState *runtime)
 {
-    runtime->emscripten_count_args_function = _PyEM_GetCountArgsPtr();
+  runtime->emscripten_trampoline = _PyEM_GetTrampolinePtr();
 }
 
 // We have to be careful to work correctly with memory snapshots. Even if we are
@@ -196,23 +101,16 @@ _PyEM_TrampolineCall(PyCFunctionWithKeywords func,
                      PyObject* args,
                      PyObject* kw)
 {
-    CountArgsFunc count_args = _PyRuntime.emscripten_count_args_function;
-    if (count_args == 0) {
-        return _PyEM_TrampolineCall_JS(func, self, args, kw);
-    }
-    switch (count_args(func)) {
-        case 0:
-            return ((zero_arg)func)();
-        case 1:
-            return ((one_arg)func)(self);
-        case 2:
-            return ((two_arg)func)(self, args);
-        case 3:
-            return ((three_arg)func)(self, args, kw);
-        default:
-            PyErr_SetString(PyExc_SystemError, "Handler takes too many arguments");
-            return NULL;
-    }
+  TrampolineFunc trampoline = _PyRuntime.emscripten_trampoline;
+  if (trampoline == 0) {
+    return _PyEM_TrampolineCall_JS(func, self, args, kw);
+  }
+  int success = 0;
+  PyObject *result = trampoline(&success, func, self, args, kw);
+  if (!success) {
+    PyErr_SetString(PyExc_SystemError, "Handler takes too many arguments");
+  }
+  return result;
 }
 
 #endif
diff --git a/configure b/configure
diff --git a/configure.ac b/configure.ac
@@ -2336,7 +2336,7 @@ AS_CASE([$ac_sys_system],
     dnl Include file system support
     AS_VAR_APPEND([LINKFORSHARED], [" -sFORCE_FILESYSTEM -lidbfs.js -lnodefs.js -lproxyfs.js -lworkerfs.js"])
     AS_VAR_APPEND([LINKFORSHARED], [" -sEXPORTED_RUNTIME_METHODS=FS,callMain,ENV,HEAPU32,TTY"])
-    AS_VAR_APPEND([LINKFORSHARED], [" -sEXPORTED_FUNCTIONS=_main,_Py_Version,__PyRuntime,__PyEM_EMSCRIPTEN_COUNT_ARGS_OFFSET,_PyGILState_GetThisThreadState,__Py_DumpTraceback"])
+    AS_VAR_APPEND([LINKFORSHARED], [" -sEXPORTED_FUNCTIONS=_main,_Py_Version,__PyRuntime,__PyEM_EMSCRIPTEN_TRAMPOLINE_OFFSET,_PyGILState_GetThisThreadState,__Py_DumpTraceback"])
     AS_VAR_APPEND([LINKFORSHARED], [" -sSTACK_SIZE=5MB"])
     dnl Avoid bugs in JS fallback string decoding path
     AS_VAR_APPEND([LINKFORSHARED], [" -sTEXTDECODER=2"])
