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.

1 files changed, 18 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) {