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.

10 files changed, 67 insertions, 47 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;
diff --git a/test/dot_s/bcp-1.wast b/test/dot_s/bcp-1.wast
index 4e3fa51ea..659f9c7b7 100644
--- a/test/dot_s/bcp-1.wast
+++ b/test/dot_s/bcp-1.wast
@@ -11,10 +11,10 @@
   (export "memory" memory)
   (type $FUNCSIG$i (func (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
-  (type $FUNCSIG$vi (func (param i32)))
   (type $FUNCSIG$v (func))
-  (import $exit "env" "exit" (param i32))
+  (type $FUNCSIG$vi (func (param i32)))
   (import $abort "env" "abort")
+  (import $exit "env" "exit" (param i32))
   (export "bad0" $bad0)
   (export "bad1" $bad1)
   (export "bad2" $bad2)
diff --git a/test/llvm_autogenerated/call.wast b/test/llvm_autogenerated/call.wast
index d7142a446..abaa2454f 100644
--- a/test/llvm_autogenerated/call.wast
+++ b/test/llvm_autogenerated/call.wast
@@ -3,20 +3,20 @@
     (segment 4 "\10\04\00\00")
   )
   (export "memory" memory)
+  (type $FUNCSIG$v (func))
   (type $FUNCSIG$i (func (result i32)))
-  (type $FUNCSIG$j (func (result i64)))
-  (type $FUNCSIG$f (func (result f32)))
   (type $FUNCSIG$d (func (result f64)))
-  (type $FUNCSIG$v (func))
-  (type $FUNCSIG$ii (func (param i32) (result i32)))
+  (type $FUNCSIG$f (func (result f32)))
   (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
+  (type $FUNCSIG$ii (func (param i32) (result i32)))
+  (type $FUNCSIG$j (func (result i64)))
+  (import $double_nullary "env" "double_nullary" (result f64))
+  (import $float_nullary "env" "float_nullary" (result f32))
+  (import $i32_binary "env" "i32_binary" (param i32 i32) (result i32))
   (import $i32_nullary "env" "i32_nullary" (result i32))
+  (import $i32_unary "env" "i32_unary" (param i32) (result i32))
   (import $i64_nullary "env" "i64_nullary" (result i64))
-  (import $float_nullary "env" "float_nullary" (result f32))
-  (import $double_nullary "env" "double_nullary" (result f64))
   (import $void_nullary "env" "void_nullary")
-  (import $i32_unary "env" "i32_unary" (param i32) (result i32))
-  (import $i32_binary "env" "i32_binary" (param i32 i32) (result i32))
   (export "call_i32_nullary" $call_i32_nullary)
   (export "call_i64_nullary" $call_i64_nullary)
   (export "call_float_nullary" $call_float_nullary)
diff --git a/test/llvm_autogenerated/cfg-stackify.wast b/test/llvm_autogenerated/cfg-stackify.wast
index 36c3e389b..18475333c 100644
--- a/test/llvm_autogenerated/cfg-stackify.wast
+++ b/test/llvm_autogenerated/cfg-stackify.wast
@@ -3,11 +3,11 @@
     (segment 4 "\10\04\00\00")
   )
   (export "memory" memory)
-  (type $FUNCSIG$v (func))
   (type $FUNCSIG$i (func (result i32)))
-  (import $something "env" "something")
-  (import $bar "env" "bar")
+  (type $FUNCSIG$v (func))
   (import $a "env" "a" (result i32))
+  (import $bar "env" "bar")
+  (import $something "env" "something")
   (export "test0" $test0)
   (export "test1" $test1)
   (export "test2" $test2)
diff --git a/test/llvm_autogenerated/frem.wast b/test/llvm_autogenerated/frem.wast
index 4b15f3a14..2d7d4be4a 100644
--- a/test/llvm_autogenerated/frem.wast
+++ b/test/llvm_autogenerated/frem.wast
@@ -3,10 +3,10 @@
     (segment 4 "\10\04\00\00")
   )
   (export "memory" memory)
-  (type $FUNCSIG$fff (func (param f32 f32) (result f32)))
   (type $FUNCSIG$ddd (func (param f64 f64) (result f64)))
-  (import $fmodf "env" "fmodf" (param f32 f32) (result f32))
+  (type $FUNCSIG$fff (func (param f32 f32) (result f32)))
   (import $fmod "env" "fmod" (param f64 f64) (result f64))
+  (import $fmodf "env" "fmodf" (param f32 f32) (result f32))
   (export "frem32" $frem32)
   (export "frem64" $frem64)
   (func $frem32 (param $$0 f32) (param $$1 f32) (result f32)
diff --git a/test/llvm_autogenerated/legalize.wast b/test/llvm_autogenerated/legalize.wast
index 1ecc46fca..0da123266 100644
--- a/test/llvm_autogenerated/legalize.wast
+++ b/test/llvm_autogenerated/legalize.wast
@@ -4,8 +4,8 @@
   )
   (export "memory" memory)
   (type $FUNCSIG$vijjj (func (param i32 i64 i64 i64)))
-  (import $__lshrti3 "env" "__lshrti3" (param i32 i64 i64 i64))
   (import $__ashlti3 "env" "__ashlti3" (param i32 i64 i64 i64))
+  (import $__lshrti3 "env" "__lshrti3" (param i32 i64 i64 i64))
   (export "shl_i3" $shl_i3)
   (export "shl_i53" $shl_i53)
   (export "sext_in_reg_i32_i64" $sext_in_reg_i32_i64)
diff --git a/test/llvm_autogenerated/mem-intrinsics.wast b/test/llvm_autogenerated/mem-intrinsics.wast
index d5456299a..f5f36ab69 100644
--- a/test/llvm_autogenerated/mem-intrinsics.wast
+++ b/test/llvm_autogenerated/mem-intrinsics.wast
@@ -3,12 +3,12 @@
     (segment 4 "\10\04\00\00")
   )
   (export "memory" memory)
-  (type $FUNCSIG$iiii (func (param i32 i32 i32) (result i32)))
   (type $FUNCSIG$i (func (result i32)))
+  (type $FUNCSIG$iiii (func (param i32 i32 i32) (result i32)))
+  (import $def "env" "def" (result i32))
   (import $memcpy "env" "memcpy" (param i32 i32 i32) (result i32))
   (import $memmove "env" "memmove" (param i32 i32 i32) (result i32))
   (import $memset "env" "memset" (param i32 i32 i32) (result i32))
-  (import $def "env" "def" (result i32))
   (export "copy_yes" $copy_yes)
   (export "copy_no" $copy_no)
   (export "move_yes" $move_yes)
diff --git a/test/llvm_autogenerated/reg-stackify.wast b/test/llvm_autogenerated/reg-stackify.wast
index c35651002..14a8d8111 100644
--- a/test/llvm_autogenerated/reg-stackify.wast
+++ b/test/llvm_autogenerated/reg-stackify.wast
@@ -3,19 +3,19 @@
     (segment 4 "\10\04\00\00")
   )
   (export "memory" memory)
-  (type $FUNCSIG$v (func))
-  (type $FUNCSIG$vi (func (param i32)))
-  (type $FUNCSIG$vii (func (param i32 i32)))
   (type $FUNCSIG$i (func (result i32)))
   (type $FUNCSIG$ii (func (param i32) (result i32)))
+  (type $FUNCSIG$v (func))
+  (type $FUNCSIG$vii (func (param i32 i32)))
+  (type $FUNCSIG$vi (func (param i32)))
+  (import $blue "env" "blue" (result i32))
+  (import $callee "env" "callee" (param i32) (result i32))
   (import $evoke_side_effects "env" "evoke_side_effects")
+  (import $green "env" "green" (result i32))
+  (import $red "env" "red" (result i32))
+  (import $use_2 "env" "use_2" (param i32 i32))
   (import $use_a "env" "use_a" (param i32))
   (import $use_b "env" "use_b" (param i32))
-  (import $use_2 "env" "use_2" (param i32 i32))
-  (import $red "env" "red" (result i32))
-  (import $green "env" "green" (result i32))
-  (import $blue "env" "blue" (result i32))
-  (import $callee "env" "callee" (param i32) (result i32))
   (export "no0" $no0)
   (export "no1" $no1)
   (export "yes0" $yes0)