gh-128627: Use __builtin_wasm_test_function_pointer_signature for Emscripten trampoline

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. */

Makefile.pre.in

@@ -3106,6 +3106,12 @@ 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
+	# emcc has a path that ends with emsdk/upstream/emscripten/emcc, we're looking for emsdk/upstream/bin/clang.
+	$$(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 \

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 _PyEM_trampoline_inner_wasm[] = {
+#embed "Python/emscripten_trampoline_inner.wasm"
+};
+
+EMSCRIPTEN_KEEPALIVE const int _PyEM_trampoline_inner_wasm_length =
+    sizeof(trampoline_inner_wasm);
 
 // Offset of emscripten_count_args_function in _PyRuntimeState. There's a couple
 // of alternatives:
@@ -18,59 +27,14 @@ 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);
+EMSCRIPTEN_KEEPALIVE const int _PyEM_EMSCRIPTEN_TRAMPOLINE_OFFSET =
+    offsetof(_PyRuntimeState, emscripten_trampoline);
 
-EM_JS(CountArgsFunc, _PyEM_GetCountArgsPtr, (), {
+EM_JS(TrampolineFunc, _PyEM_GetTrampolinePtr, (void), {
     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() {
+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.
@@ -84,74 +48,13 @@ function getPyEMCountArgsPtr() {
     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
-    ]);
+    const code = HEAP8.subarray(
+        __PyEM_trampoline_inner_wasm,
+        __PyEM_trampoline_inner_wasm + HEAP32[__PyEM_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);
+        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.
@@ -160,17 +63,17 @@ function getPyEMCountArgsPtr() {
 }
 
 addOnPreRun(() => {
-    const ptr = getPyEMCountArgsPtr();
+    const ptr = getPyEMTrampolinePtr();
     Module._PyEM_CountArgsPtr = ptr;
-    const offset = HEAP32[__PyEM_EMSCRIPTEN_COUNT_ARGS_OFFSET / 4];
+    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 +99,16 @@ _PyEM_TrampolineCall(PyCFunctionWithKeywords func,
                      PyObject* args,
                      PyObject* kw)
 {
-    CountArgsFunc count_args = _PyRuntime.emscripten_count_args_function;
-    if (count_args == 0) {
+    TrampolineFunc trampoline = _PyRuntime.emscripten_trampoline;
+    if (trampoline == 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;
+    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

Python/emscripten_trampoline_inner.c

@@ -0,0 +1,27 @@
+typedef void PyObject;
+
+typedef PyObject* (*zero_arg)(void);
+typedef PyObject* (*one_arg)(PyObject*);
+typedef PyObject* (*two_arg)(PyObject*, PyObject*);
+typedef PyObject* (*three_arg)(PyObject*, PyObject*, PyObject*);
+
+#define TRY_RETURN_CALL(ty, args...) \
+  if (__builtin_wasm_test_function_pointer_signature((ty)func)) { \
+    return ((ty)func)(args); \
+  }
+
+__attribute__((export_name("trampoline_call"))) PyObject*
+trampoline_call(int* success,
+                void* func,
+                PyObject* self,
+                PyObject* args,
+                PyObject* kw)
+{
+  *success = 1;
+  TRY_RETURN_CALL(three_arg, self, args, kw);
+  TRY_RETURN_CALL(two_arg, self, args);
+  TRY_RETURN_CALL(one_arg, self);
+  TRY_RETURN_CALL(zero_arg);
+  *success = 0;
+  return 0;
+}

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"])