Defer creation of CallImports to link time (#395)

s2wasm currently creates a Call AST node if the target is implemented in the current object (thus far assumed to be the final executable) and a CallImport node if not. In preparation for adding additional objects to the link before layout time, we make only Call nodes until link time, and then convert them to CallImport if they are undefined at that time.

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;