gh-128627: Emscripten: Use wasm-gc based call adaptor if available (#128628)

Replaces the trampoline mechanism in Emscripten with an implementation that uses a
recently added feature of wasm-gc instead of JS type reflection, when that feature is
available.

1 files changed, 177 insertions, 52 deletions

diff --git a/Python/emscripten_trampoline.c b/Python/emscripten_trampoline.c
index 960c6b4..2f9648a 100644
--- a/Python/emscripten_trampoline.c
+++ b/Python/emscripten_trampoline.c
@@ -1,79 +1,205 @@
 #if defined(PY_CALL_TRAMPOLINE)
 
-#include <emscripten.h>             // EM_JS
+#include <emscripten.h>             // EM_JS, EM_JS_DEPS
 #include <Python.h>
 #include "pycore_runtime.h"         // _PyRuntime
 
+typedef int (*CountArgsFunc)(PyCFunctionWithKeywords func);
 
-/**
- * This is the GoogleChromeLabs approved way to feature detect type-reflection:
- * https://github.com/GoogleChromeLabs/wasm-feature-detect/blob/main/src/detectors/type-reflection/index.js
- */
-EM_JS(int, _PyEM_detect_type_reflection, (), {
-    if (!("Function" in WebAssembly)) {
-        return false;
-    }
-    if (WebAssembly.Function.type) {
-        // Node v20
-        Module.PyEM_CountArgs = (func) => WebAssembly.Function.type(wasmTable.get(func)).parameters.length;
-    } else {
-        // Node >= 22, v8-based browsers
-        Module.PyEM_CountArgs = (func) => wasmTable.get(func).type().parameters.length;
+// Offset of emscripten_count_args_function in _PyRuntimeState. There's a couple
+// of alternatives:
+// 1. Just make emscripten_count_args_function a real C global variable instead
+//    of a field of _PyRuntimeState. This would violate our rule against mutable
+//    globals.
+// 2. #define a preprocessor constant equal to a hard coded number and make a
+//    _Static_assert(offsetof(_PyRuntimeState, emscripten_count_args_function)
+//    == OURCONSTANT) This has the disadvantage that we have to update the hard
+//    coded constant when _PyRuntimeState changes
+//
+// 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 module to the binary, my approach is to find and
+// replace "ref.test $type" -> "drop i32.const n" on the source text. This
+// results in the bytes "0x1a, 0x41, n" where we need the bytes "0xfb, 0x14, n"
+// so doing a find and replace on the output from "0x1a, 0x41" -> "0xfb, 0x14"
+// gets us the output we need.
+//
+// (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 i32) (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
+//         (block $b (type $blocktype)
+//             i32.eqz
+//             br_if $b
+//             i32.const 3
+//             return
+//         )
+//         local.get $fref
+//         ref.test $type2
+//         (block $b (type $blocktype)
+//             i32.eqz
+//             br_if $b
+//             i32.const 2
+//             return
+//         )
+//         local.get $fref
+//         ref.test $type1
+//         (block $b (type $blocktype)
+//             i32.eqz
+//             br_if $b
+//             i32.const 1
+//             return
+//         )
+//         local.get $fref
+//         ref.test $type0
+//         (block $b (type $blocktype)
+//             i32.eqz
+//             br_if $b
+//             i32.const 0
+//             return
+//         )
+//         i32.const -1
+//     )
+// )
+addOnPreRun(() => {
+    // Try to initialize countArgsFunc
+    const code = new Uint8Array([
+        0x00, 0x61, 0x73, 0x6d, // \0asm magic number
+        0x01, 0x00, 0x00, 0x00, // version 1
+        0x01, 0x1b, // Type section, body is 0x1b 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, 0x01, 0x7f, 0x00,                         // (type $blocktype (func (param i32) (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, 0x44,  // Code section,
+            0x01, 0x42, // one entry of length 50
+            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
+            0x02, 0x04,       // block $b (type $blocktype)
+                0x45,         //   i32.eqz
+                0x0d, 0x00,   //   br_if $b
+                0x41, 0x03,   //   i32.const 3
+                0x0f,         //   return
+            0x0b,             // end block
+
+            0x20, 0x01,       // local.get $fref
+            0xfb, 0x14, 0x02, // ref.test $type2
+            0x02, 0x04,       // block $b (type $blocktype)
+                0x45,         //   i32.eqz
+                0x0d, 0x00,   //   br_if $b
+                0x41, 0x02,   //   i32.const 2
+                0x0f,         //   return
+            0x0b,             // end block
+
+            0x20, 0x01,       // local.get $fref
+            0xfb, 0x14, 0x01, // ref.test $type1
+            0x02, 0x04,       // block $b (type $blocktype)
+                0x45,         //   i32.eqz
+                0x0d, 0x00,   //   br_if $b
+                0x41, 0x01,   //   i32.const 1
+                0x0f,         //   return
+            0x0b,             // end block
+
+            0x20, 0x01,       // local.get $fref
+            0xfb, 0x14, 0x00, // ref.test $type0
+            0x02, 0x04,       // block $b (type $blocktype)
+                0x45,         //   i32.eqz
+                0x0d, 0x00,   //   br_if $b
+                0x41, 0x00,   //   i32.const 0
+                0x0f,         //   return
+            0x0b,             // end block
+
+            0x41, 0x7f,       // i32.const -1
+            0x0b // end function
+    ]);
+    let ptr = 0;
+    try {
+        const mod = new WebAssembly.Module(code);
+        const inst = new WebAssembly.Instance(mod, {e: {t: wasmTable}});
+        ptr = addFunction(inst.exports.f);
+    } catch(e) {
+        // If something goes wrong, we'll null out _PyEM_CountFuncParams and fall
+        // back to the JS trampoline.
     }
-    return true;
+    Module._PyEM_CountArgsPtr = ptr;
+    const offset = HEAP32[__PyEM_EMSCRIPTEN_COUNT_ARGS_OFFSET/4];
+    HEAP32[__PyRuntime/4 + offset] = ptr;
 });
+);
 
 void
 _Py_EmscriptenTrampoline_Init(_PyRuntimeState *runtime)
 {
-    runtime->wasm_type_reflection_available = _PyEM_detect_type_reflection();
+    runtime->emscripten_count_args_function = _PyEM_GetCountArgsPtr();
 }
 
+// We have to be careful to work correctly with memory snapshots. Even if we are
+// loading a memory snapshot, we need to perform the JS initialization work.
+// That means we can't call the initialization code from C. Instead, we export
+// this function pointer to JS and then fill it in a preRun function which runs
+// unconditionally.
 /**
  * Backwards compatible trampoline works with all JS runtimes
  */
-EM_JS(PyObject*,
-_PyEM_TrampolineCall_JavaScript, (PyCFunctionWithKeywords func,
-                                  PyObject *arg1,
-                                  PyObject *arg2,
-                                  PyObject *arg3),
-{
+EM_JS(PyObject*, _PyEM_TrampolineCall_JS, (PyCFunctionWithKeywords func, PyObject *arg1, PyObject *arg2, PyObject *arg3), {
     return wasmTable.get(func)(arg1, arg2, arg3);
-}
-);
-
-/**
- * In runtimes with WebAssembly type reflection, count the number of parameters
- * and cast to the appropriate signature
- */
-EM_JS(int, _PyEM_CountFuncParams, (PyCFunctionWithKeywords func),
-{
-    let n = _PyEM_CountFuncParams.cache.get(func);
-
-    if (n !== undefined) {
-        return n;
-    }
-    n = Module.PyEM_CountArgs(func);
-    _PyEM_CountFuncParams.cache.set(func, n);
-    return n;
-}
-_PyEM_CountFuncParams.cache = new Map();
-)
-
+});
 
 typedef PyObject* (*zero_arg)(void);
 typedef PyObject* (*one_arg)(PyObject*);
 typedef PyObject* (*two_arg)(PyObject*, PyObject*);
 typedef PyObject* (*three_arg)(PyObject*, PyObject*, PyObject*);
 
-
 PyObject*
-_PyEM_TrampolineCall_Reflection(PyCFunctionWithKeywords func,
-                                PyObject* self,
-                                PyObject* args,
-                                PyObject* kw)
+_PyEM_TrampolineCall(PyCFunctionWithKeywords func,
+                     PyObject* self,
+                     PyObject* args,
+                     PyObject* kw)
 {
-    switch (_PyEM_CountFuncParams(func)) {
+    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:
@@ -83,8 +209,7 @@ _PyEM_TrampolineCall_Reflection(PyCFunctionWithKeywords func,
         case 3:
             return ((three_arg)func)(self, args, kw);
         default:
-            PyErr_SetString(PyExc_SystemError,
-                            "Handler takes too many arguments");
+            PyErr_SetString(PyExc_SystemError, "Handler takes too many arguments");
             return NULL;
     }
 }