JSPI - Replace function table references with JSPI'ed wrapper. (#5519)

This makes it possible to get the JSPI'ed version of the function from the
function table.

2 files changed, 52 insertions, 0 deletions

diff --git a/src/passes/JSPI.cpp b/src/passes/JSPI.cpp
index 01532bb38..a3b544bca 100644
--- a/src/passes/JSPI.cpp
+++ b/src/passes/JSPI.cpp
@@ -140,6 +140,24 @@ struct JSPI : public Pass {
       }
     }
 
+    // Replace any references to the original exports that are in the elements.
+    for (auto& segment : module->elementSegments) {
+      if (!segment->type.isFunction()) {
+        continue;
+      }
+      for (Index i = 0; i < segment->data.size(); i++) {
+        if (auto* get = segment->data[i]->dynCast<RefFunc>()) {
+          auto iter = wrappedExports.find(get->func);
+          if (iter == wrappedExports.end()) {
+            continue;
+          }
+          auto* replacementRef = builder.makeRefFunc(
+            iter->second, module->getFunction(iter->second)->type);
+          segment->data[i] = replacementRef;
+        }
+      }
+    }
+
     // Avoid iterator invalidation later.
     std::vector<Function*> originalFunctions;
     for (auto& func : module->functions) {
diff --git a/test/lit/passes/jspi-table.wast b/test/lit/passes/jspi-table.wast
new file mode 100644
index 000000000..1963d5cb0
--- /dev/null
+++ b/test/lit/passes/jspi-table.wast
@@ -0,0 +1,34 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt %s --jspi -all -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $f64_=>_i32 (func (param f64) (result i32)))
+
+  ;; CHECK:      (type $externref_f64_=>_i32 (func (param externref f64) (result i32)))
+
+  ;; CHECK:      (global $suspender (mut externref) (ref.null noextern))
+
+  ;; CHECK:      (table $0 1 1 funcref)
+  (table $0 1 1 funcref)
+  (elem (i32.const 1) func $update_state)
+  ;; CHECK:      (elem (i32.const 1) $export$update_state)
+
+  ;; CHECK:      (export "update_state" (func $export$update_state))
+  (export "update_state" (func $update_state))
+
+  ;; CHECK:      (func $update_state (type $f64_=>_i32) (param $param f64) (result i32)
+  ;; CHECK-NEXT:  (i32.const 42)
+  ;; CHECK-NEXT: )
+  (func $update_state (param $param f64) (result i32)
+    (i32.const 42)
+  )
+)
+;; CHECK:      (func $export$update_state (type $externref_f64_=>_i32) (param $susp externref) (param $param f64) (result i32)
+;; CHECK-NEXT:  (global.set $suspender
+;; CHECK-NEXT:   (local.get $susp)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $update_state
+;; CHECK-NEXT:   (local.get $param)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )