3 files changed, 41 insertions, 21 deletions

diff --git a/src/s2wasm.h b/src/s2wasm.h
index f70246d4f..94684facf 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -696,19 +696,15 @@ class S2WasmBuilder {
 
       } else {
         // non-indirect call
-        CallBase* curr;
         Name assign = getAssign();
         Name target = linkerObj.resolveAlias(cleanFunction(getCommaSeparated()));
-        if (linkerObj.isFunctionImplemented(target)) {
-          auto specific = allocator.alloc<Call>();
-          specific->target = target;
-          curr = specific;
-        } else {
-          auto specific = allocator.alloc<CallImport>();
-          specific->target = target;
-          curr = specific;
-        }
+
+        Call* curr = allocator.alloc<Call>();
+        curr->target = target;
         curr->type = type;
+        if (!linkerObj.isFunctionImplemented(target)) {
+          linkerObj.addUndefinedFunctionCall(curr);
+        }
         skipWhitespace();
         if (*s == ',') {
           skipComma();
@@ -719,16 +715,6 @@ class S2WasmBuilder {
           }
         }
         setOutput(curr, assign);
-        if (curr->is<CallImport>()) {
-          auto target = curr->cast<CallImport>()->target;
-          if (!wasm.checkImport(target)) {
-            auto import = allocator.alloc<Import>();
-            import->name = import->base = target;
-            import->module = ENV;
-            import->type = ensureFunctionType(getSig(curr), &wasm, allocator);
-            wasm.addImport(import);
-          }
-        }
       }
     };
     auto handleTyped = [&](WasmType type) {
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 2e5787670..4fb7ea4ae 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -49,13 +49,40 @@ void Linker::placeStackPointer(size_t stackAllocation) {
 }
 
 void Linker::layout() {
+  // Convert calls to undefined functions to call_imports
+  for (const auto& f : out.undefinedFunctions) {
+    Name target = f.first;
+    // Create an import for the target if necessary.
+    if (!out.wasm.checkImport(target)) {
+      auto import = out.wasm.allocator.alloc<Import>();
+      import->name = import->base = target;
+      import->module = ENV;
+      import->type = ensureFunctionType(getSig(*f.second.begin()), &out.wasm,
+                                        out.wasm.allocator);
+      out.wasm.addImport(import);
+    }
+    // Change each call. The target is the same since it's still the name.
+    // Delete and re-allocate the Expression as CallImport to avoid undefined
+    // behavior.
+    static_assert(sizeof(Call) >= sizeof(CallImport),
+                  "Cannot reallocate a CallImport in a Call arena slot");
+    for (auto* call : f.second) {
+      Call callCopy = std::move(*call);
+      call->~Call();
+      CallImport* newCall = new (call) CallImport;
+      newCall->type = callCopy.type;
+      newCall->operands = std::move(callCopy.operands);
+      newCall->target = target;
+    }
+  }
+
   // Allocate all user statics
   for (const auto& obj : out.staticObjects) {
     allocateStatic(obj.allocSize, obj.alignment, obj.name);
   }
 
   // Update the segments with their addresses now that they have been allocated.
-  for (auto& seg : out.segments) {
+  for (const auto& seg : out.segments) {
     size_t address = staticAddresses[seg.first];
     out.wasm.memory.segments[seg.second].offset = address;
     segmentsByAddress[address] = seg.second;
diff --git a/src/wasm-linker.h b/src/wasm-linker.h
index 023cf198c..e72f862f1 100644
--- a/src/wasm-linker.h
+++ b/src/wasm-linker.h
@@ -91,6 +91,10 @@ class LinkerObject {
     assert(implementedFunctions.count(name));
   }
 
+  void addUndefinedFunctionCall(Call* call) {
+    undefinedFunctions[call->target].push_back(call);
+  }
+
   bool isEmpty() {
     return wasm.functions.empty();
   }
@@ -116,6 +120,9 @@ class LinkerObject {
   std::set<Name> implementedFunctions;
   std::unordered_map<cashew::IString, Name> aliasedFunctions;
 
+  using CallList = std::vector<Call*>;
+  std::map<Name, CallList> undefinedFunctions;
+
   std::map<Name, size_t> segments; // name => segment index (in wasm module)
 
   std::vector<Name> initializerFunctions;