wasm-merge: Fix locals in merged start (#5837)

Start functions can have locals, which we previously ignored as we just
concatenated the bodies together. This makes us copy the second start
and call that, keeping them separate (the optimizer can then inline, if that
makes sense).

Fixes #5835

4 files changed, 45 insertions, 15 deletions

diff --git a/src/tools/wasm-merge.cpp b/src/tools/wasm-merge.cpp
index 9fe97669b..8df93196f 100644
--- a/src/tools/wasm-merge.cpp
+++ b/src/tools/wasm-merge.cpp
@@ -362,19 +362,21 @@ void copyModuleContents(Module& input, Name inputName) {
       // No previous start; just refer to the new one.
       merged.start = input.start;
     } else {
-      // Merge them, keeping the order. Note that we need to create a new
-      // function as there may be other references.
+      // Merge them, keeping the order. We copy both functions to avoid issues
+      // with other references to them, and just call the second one, leaving
+      // inlining to the optimizer if that makes sense to do.
+      auto copiedOldName =
+        Names::getValidFunctionName(merged, "merged.start.old");
+      auto copiedNewName =
+        Names::getValidFunctionName(merged, "merged.start.new");
+      auto* copiedOld = ModuleUtils::copyFunction(
+        merged.getFunction(merged.start), merged, copiedOldName);
+      ModuleUtils::copyFunction(
+        merged.getFunction(input.start), merged, copiedNewName);
       Builder builder(merged);
-      auto mergedName = Names::getValidFunctionName(merged, "merged.start");
-      auto* oldStart = merged.getFunction(merged.start);
-      auto* oldStartBody = ExpressionManipulator::copy(oldStart->body, merged);
-      auto* newStart = merged.getFunction(input.start);
-      auto* newStartBody = ExpressionManipulator::copy(newStart->body, merged);
-      auto* mergedBody = builder.makeSequence(oldStartBody, newStartBody);
-      auto mergedFunc = builder.makeFunction(
-        mergedName, Signature{Type::none, Type::none}, {}, mergedBody);
-      merged.addFunction(std::move(mergedFunc));
-      merged.start = mergedName;
+      copiedOld->body = builder.makeSequence(
+        copiedOld->body, builder.makeCall(copiedNewName, {}, Type::none));
+      merged.start = copiedOldName;
     }
   }
 
diff --git a/test/lit/merge/start3.wat b/test/lit/merge/start3.wat
index ab4c555ab..9918d0164 100644
--- a/test/lit/merge/start3.wat
+++ b/test/lit/merge/start3.wat
@@ -13,9 +13,13 @@
 
 ;; CHECK:      (export "user" (func $user))
 
-;; CHECK:      (start $merged.start)
+;; CHECK:      (start $merged.start.old)
 
 ;; CHECK:      (func $start (type $none_=>_none)
+;; CHECK-NEXT:  (local $x i32)
+;; CHECK-NEXT:  (drop
+;; CHECK-NEXT:   (local.get $x)
+;; CHECK-NEXT:  )
 ;; CHECK-NEXT:  (drop
 ;; CHECK-NEXT:   (i32.const 1)
 ;; CHECK-NEXT:  )
@@ -26,9 +30,23 @@
 ;; CHECK-NEXT:  (call $start)
 ;; CHECK-NEXT: )
 
-;; CHECK:      (func $merged.start (type $none_=>_none)
+;; CHECK:      (func $merged.start.old (type $none_=>_none)
+;; CHECK-NEXT:  (local $x i32)
+;; CHECK-NEXT:  (block
+;; CHECK-NEXT:   (drop
+;; CHECK-NEXT:    (local.get $x)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (drop
+;; CHECK-NEXT:    (i32.const 1)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $merged.start.new)
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $merged.start.new (type $none_=>_none)
+;; CHECK-NEXT:  (local $x f64)
 ;; CHECK-NEXT:  (drop
-;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:   (local.get $x)
 ;; CHECK-NEXT:  )
 ;; CHECK-NEXT:  (drop
 ;; CHECK-NEXT:   (i32.const 2)
diff --git a/test/lit/merge/start3.wat.second b/test/lit/merge/start3.wat.second
index ec66c2f5b..39090b56d 100644
--- a/test/lit/merge/start3.wat.second
+++ b/test/lit/merge/start3.wat.second
@@ -2,6 +2,12 @@
   (start $start)
 
   (func $start (export "start")
+    ;; Test that locals are handled properly. This function has an i32 local,
+    ;; and the other start has an f64 with the same name.
+    (local $x i32)
+    (drop
+      (local.get $x)
+    )
     (drop
       (i32.const 1)
     )
diff --git a/test/lit/merge/start3.wat.third b/test/lit/merge/start3.wat.third
index 508ffdb2a..e4888b66e 100644
--- a/test/lit/merge/start3.wat.third
+++ b/test/lit/merge/start3.wat.third
@@ -2,6 +2,10 @@
   (start $start)
 
   (func $start
+    (local $x f64)
+    (drop
+      (local.get $x)
+    )
     (drop
       (i32.const 2)
     )