From 20fe882b47c4b2570c70b2f9d82189c5b2144d03 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Tue, 14 Nov 2023 19:09:49 -0800
Subject: [PATCH] SignatureRefining: Notice LUB requirements of intrinsic calls
 (#6122)

call.without.effects implies a call to the function reference in the last parameter,
so the values sent in the other parameters must be taken into account when
computing LUBs for refining arguments, otherwise we might refine so much that
the intrinsic call no longer validates.
---
 src/passes/SignatureRefining.cpp        | 26 +++++++++++
 test/lit/passes/signature-refining.wast | 57 +++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/src/passes/SignatureRefining.cpp b/src/passes/SignatureRefining.cpp
index 24d120dda88..ae124dc79a0 100644
--- a/src/passes/SignatureRefining.cpp
+++ b/src/passes/SignatureRefining.cpp
@@ -73,6 +73,15 @@ struct SignatureRefining : public Pass {
       std::vector<Call*> calls;
       std::vector<CallRef*> callRefs;
 
+      // Additional calls to take into account. We store intrinsic calls here,
+      // as they must appear twice: call.without.effects is both a normal call
+      // and also takes a final parameter that is a function reference that is
+      // called, and so two signatures are relevant for it. For the latter, we
+      // add the call as an "extra call" (which is an unusual call, as it has an
+      // extra parameter at the end, the function reference, compared to what we
+      // expect for the signature being called).
+      std::vector<Call*> extraCalls;
+
       // A possibly improved LUB for the results.
       LUBFinder resultsLUB;
 
@@ -107,6 +116,17 @@ struct SignatureRefining : public Pass {
       // called.
       for (auto* call : info.calls) {
         allInfo[module->getFunction(call->target)->type].calls.push_back(call);
+
+        // For call.without.effects, we also add the effective function being
+        // called as well. The final operand is the function reference being
+        // called, which defines that type.
+        if (Intrinsics(*module).isCallWithoutEffects(call)) {
+          auto targetType = call->operands.back()->type;
+          if (!targetType.isRef()) {
+            continue;
+          }
+          allInfo[targetType.getHeapType()].extraCalls.push_back(call);
+        }
       }
 
       // For indirect calls, add each call_ref to the type the call_ref uses.
@@ -187,6 +207,12 @@ struct SignatureRefining : public Pass {
       for (auto* callRef : info.callRefs) {
         updateLUBs(callRef->operands);
       }
+      for (auto* call : info.extraCalls) {
+        // Note that these intrinsic calls have an extra function reference
+        // param at the end, but updateLUBs looks at |numParams| only, so it
+        // considers just the relevant parameters.
+        updateLUBs(call->operands);
+      }
 
       // Find the final LUBs, and see if we found an improvement.
       std::vector<Type> newParamsTypes;
diff --git a/test/lit/passes/signature-refining.wast b/test/lit/passes/signature-refining.wast
index c2a505e5151..95007b1c4af 100644
--- a/test/lit/passes/signature-refining.wast
+++ b/test/lit/passes/signature-refining.wast
@@ -985,3 +985,60 @@
   (struct.new $C) ;; this will allow this function's result to be refined to $C
  )
 )
+
+;; Test we consider call.without.effects when deciding what to refine. $A has
+;; two subtypes, B1 and B2, and a call.without.effects sends in one while a
+;; normal call sends in the other. As a result, we cannot refine to either.
+(module
+ (rec
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $A (sub (struct )))
+  (type $A (sub (struct)))
+
+  ;; CHECK:       (type $B1 (sub $A (struct )))
+  (type $B1 (sub $A (struct)))
+
+  ;; CHECK:       (type $B2 (sub $A (struct )))
+  (type $B2 (sub $A (struct)))
+ )
+
+ ;; CHECK:      (type $3 (func (param (ref $A) funcref)))
+
+ ;; CHECK:      (type $4 (func))
+
+ ;; CHECK:      (type $5 (func (param (ref $A))))
+
+ ;; CHECK:      (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects (type $3) (param (ref $A) funcref)))
+ (import "binaryen-intrinsics" "call.without.effects" (func $no.side.effects
+   (param (ref $A))
+   (param funcref)
+ ))
+
+ ;; CHECK:      (elem declare func $target)
+
+ ;; CHECK:      (func $calls (type $4)
+ ;; CHECK-NEXT:  (call $no.side.effects
+ ;; CHECK-NEXT:   (struct.new_default $B1)
+ ;; CHECK-NEXT:   (ref.func $target)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $target
+ ;; CHECK-NEXT:   (struct.new_default $B2)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $calls
+  (call $no.side.effects
+   (struct.new $B1)
+   (ref.func $target)
+  )
+  (call $target
+   (struct.new $B2)
+  )
+ )
+
+ ;; CHECK:      (func $target (type $5) (param $x (ref $A))
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $target (param $x (ref $A))
+  ;; Because of the two calls above, this cannot be refined.
+ )
+)