From 76fcc26e5ca62b6aa0f2d07859ff5dd95e8f57e0 Mon Sep 17 00:00:00 2001
From: Brendan Dahl <brendan.dahl@gmail.com>
Date: Wed, 1 Mar 2023 11:13:19 -0800
Subject: 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.
---
 src/passes/JSPI.cpp | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

(limited to 'src')

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) {
-- 
cgit v1.2.3