SignatureRefining: Handle updates to call.without.effects (#5884)

If we refine a signature type that is used in a call.without.effects then that call's
results may need to be updated. In the IR it looks like a normal call that happens to
pass a function reference as the last param, but it actually means that we call that
function (without side effects), so we need to have the same results, and the validator
already verified that (so the new testcase here fails without this fix).

2 files changed, 166 insertions, 1 deletions

diff --git a/src/passes/SignatureRefining.cpp b/src/passes/SignatureRefining.cpp
index 51aef429b..6d2c243c2 100644
--- a/src/passes/SignatureRefining.cpp
+++ b/src/passes/SignatureRefining.cpp
@@ -30,6 +30,7 @@
 #include "ir/find_all.h"
 #include "ir/lubs.h"
 #include "ir/module-utils.h"
+#include "ir/names.h"
 #include "ir/subtypes.h"
 #include "ir/type-updating.h"
 #include "ir/utils.h"
@@ -284,9 +285,76 @@ struct SignatureRefining : public Pass {
     // Rewrite the types.
     GlobalTypeRewriter::updateSignatures(newSignatures, *module);
 
+    // Update intrinsics.
+    updateIntrinsics(module, allInfo);
+
     // TODO: we could do this only in relevant functions perhaps
     ReFinalize().run(getPassRunner(), module);
   }
+
+  template<typename HeapInfoMap>
+  void updateIntrinsics(Module* module, HeapInfoMap& map) {
+    // The call.without.effects intrinsic needs to be updated if we refine the
+    // function reference it receives. Imagine that we have this:
+    //
+    //  (call $call.without.effects
+    //    (ref.func $returns.A)
+    //  )
+    //
+    // If we refined that $returns.A function to actually return a subtype $B,
+    // then now the call.without.effects should return $B, because logically it
+    // is still a direct call to that function. (We could also defer this to
+    // later, if we relaxed validation here, but updating right now is better
+    // for followup optimizations.)
+
+    // Each time we update we create a new import with the proper type. Keep a
+    // map of them to avoid creating more than one for each type.
+    std::unordered_map<HeapType, Function*> newImports;
+
+    auto getImportWithNewResults = [&](Function* import, Type newResults) {
+      auto newType = Signature(import->getParams(), newResults);
+      if (auto iter = newImports.find(newType); iter != newImports.end()) {
+        return iter->second;
+      }
+
+      auto name = Names::getValidFunctionName(*module, import->name);
+      auto newImport =
+        module->addFunction(Builder(*module).makeFunction(name, newType, {}));
+
+      // Copy the binaryen intrinsic module.base import names.
+      newImport->module = import->module;
+      newImport->base = import->base;
+
+      newImports[newType] = newImport;
+      return newImport;
+    };
+
+    for (auto& [_, info] : map) {
+      for (auto* call : info.calls) {
+        if (Intrinsics(*module).isCallWithoutEffects(call)) {
+          auto targetType = call->operands.back()->type;
+          if (!targetType.isRef()) {
+            continue;
+          }
+          auto heapType = targetType.getHeapType();
+          if (!heapType.isSignature()) {
+            continue;
+          }
+          auto newResults = heapType.getSignature().results;
+          if (call->type == newResults) {
+            continue;
+          }
+
+          // The target was refined, so we need to update here. Create a new
+          // import of the refined type, and call that instead.
+          call->target = getImportWithNewResults(
+                           module->getFunction(call->target), newResults)
+                           ->name;
+          call->type = newResults;
+        }
+      }
+    }
+  }
 };
 
 } // anonymous namespace
diff --git a/test/lit/passes/signature-refining.wast b/test/lit/passes/signature-refining.wast
index 6c1c23291..c4e13b975 100644
--- a/test/lit/passes/signature-refining.wast
+++ b/test/lit/passes/signature-refining.wast
@@ -887,4 +887,101 @@
    )
   )
  )
-) ;; TODO
+)
+
+;; Test the call.without.effects intrinsic, which may require additional work.
+(module
+ (rec
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $A (struct ))
+  (type $A (struct))
+
+  ;; CHECK:       (type $B (sub $A (struct )))
+  (type $B (sub $A (struct)))
+
+  ;; CHECK:       (type $C (sub $B (struct )))
+  (type $C (sub $B (struct)))
+
+  ;; CHECK:       (type $return_A_2 (func (result (ref $C))))
+
+  ;; CHECK:       (type $return_A (func (result (ref $B))))
+  (type $return_A (func (result (ref null $A))))
+
+  (type $return_A_2 (func (result (ref null $A))))
+ )
+
+ ;; CHECK:       (type $none_=>_none (func))
+
+ ;; CHECK:       (type $funcref_=>_ref?|$A| (func (param funcref) (result (ref null $A))))
+
+ ;; CHECK:      (type $funcref_=>_ref|$B| (func (param funcref) (result (ref $B))))
+
+ ;; CHECK:      (type $funcref_=>_ref|$C| (func (param funcref) (result (ref $C))))
+
+ ;; CHECK:      (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects (type $funcref_=>_ref?|$A|) (param funcref) (result (ref null $A))))
+ (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects
+   (param funcref)
+   (result (ref null $A))
+ ))
+
+ ;; CHECK:      (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects_4 (type $funcref_=>_ref|$B|) (param funcref) (result (ref $B))))
+
+ ;; CHECK:      (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects_5 (type $funcref_=>_ref|$C|) (param funcref) (result (ref $C))))
+
+ ;; CHECK:      (elem declare func $other $other2)
+
+ ;; CHECK:      (func $func (type $none_=>_none)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $no.side.effects_4
+ ;; CHECK-NEXT:    (ref.func $other)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $no.side.effects_4
+ ;; CHECK-NEXT:    (ref.func $other)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $no.side.effects_5
+ ;; CHECK-NEXT:    (ref.func $other2)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $func
+  ;; After $other's result is refined, this will need to use a new import that
+  ;; has the refined result type.
+  (drop
+   (call $no.side.effects
+    (ref.func $other)
+   )
+  )
+  ;; A second call of the same one. This should call the same new import (that
+  ;; is, we shouldn't create unnecessary copies of the new imports).
+  (drop
+   (call $no.side.effects
+    (ref.func $other)
+   )
+  )
+  ;; A call of another function with a different refining, that will need
+  ;; another import.
+  (drop
+   (call $no.side.effects
+    (ref.func $other2)
+   )
+  )
+ )
+
+ ;; CHECK:      (func $other (type $return_A) (result (ref $B))
+ ;; CHECK-NEXT:  (struct.new_default $B)
+ ;; CHECK-NEXT: )
+ (func $other (type $return_A) (result (ref null $A))
+  (struct.new $B) ;; this will allow this function's result to be refined to $B
+ )
+
+ ;; CHECK:      (func $other2 (type $return_A_2) (result (ref $C))
+ ;; CHECK-NEXT:  (struct.new_default $C)
+ ;; CHECK-NEXT: )
+ (func $other2 (type $return_A_2) (result (ref null $A))
+  (struct.new $C) ;; this will allow this function's result to be refined to $C
+ )
+)