[NFC] Refactor result type LUB computation into a helper function (#4379)

4 files changed, 123 insertions, 82 deletions

diff --git a/src/ir/CMakeLists.txt b/src/ir/CMakeLists.txt
index b625cebe7..08149e15e 100644
--- a/src/ir/CMakeLists.txt
+++ b/src/ir/CMakeLists.txt
@@ -4,6 +4,7 @@ set(ir_SOURCES
   ExpressionManipulator.cpp
   eh-utils.cpp
   intrinsics.cpp
+  lubs.cpp
   names.cpp
   properties.cpp
   LocalGraph.cpp
diff --git a/src/ir/lubs.cpp b/src/ir/lubs.cpp
new file mode 100644
index 000000000..852c27677
--- /dev/null
+++ b/src/ir/lubs.cpp
@@ -0,0 +1,99 @@
+/*
+ * Copyright 2021 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ir/lubs.h"
+#include "ir/utils.h"
+#include "wasm-type.h"
+#include "wasm.h"
+
+namespace wasm {
+
+namespace LUB {
+
+LUBFinder getResultsLUB(Function* func, Module& wasm) {
+  LUBFinder lub;
+
+  if (!wasm.features.hasGC()) {
+    return lub;
+  }
+
+  Type originalType = func->getResults();
+  if (!originalType.hasRef()) {
+    // Nothing to refine.
+    return lub;
+  }
+
+  // Before we do anything, we must refinalize the function, because otherwise
+  // its body may contain a block with a forced type,
+  //
+  // (func (result X)
+  //  (block (result X)
+  //   (..content with more specific type Y..)
+  //  )
+  ReFinalize().walkFunctionInModule(func, &wasm);
+
+  lub.noteUpdatableExpression(func->body);
+  if (lub.getBestPossible() == originalType) {
+    return lub;
+  }
+
+  // Scan the body and look at the returns. First, return expressions.
+  for (auto* ret : FindAll<Return>(func->body).list) {
+    lub.noteUpdatableExpression(ret->value);
+    if (lub.getBestPossible() == originalType) {
+      return lub;
+    }
+  }
+
+  // Process return_calls and call_refs. Unlike return expressions which we
+  // just handled, these only get a type to update, not a value.
+  auto processReturnType = [&](Type type) {
+    // Return whether we still look ok to do the optimization. If this is
+    // false then we can stop here.
+    lub.note(type);
+    return lub.getBestPossible() != originalType;
+  };
+
+  for (auto* call : FindAll<Call>(func->body).list) {
+    if (call->isReturn &&
+        !processReturnType(wasm.getFunction(call->target)->getResults())) {
+      return lub;
+    }
+  }
+  for (auto* call : FindAll<CallIndirect>(func->body).list) {
+    if (call->isReturn &&
+        !processReturnType(call->heapType.getSignature().results)) {
+      return lub;
+    }
+  }
+  for (auto* call : FindAll<CallRef>(func->body).list) {
+    if (call->isReturn) {
+      auto targetType = call->target->type;
+      if (targetType == Type::unreachable) {
+        continue;
+      }
+      if (!processReturnType(targetType.getHeapType().getSignature().results)) {
+        return lub;
+      }
+    }
+  }
+
+  return lub;
+}
+
+} // namespace LUB
+
+} // namespace wasm
diff --git a/src/ir/lubs.h b/src/ir/lubs.h
index e03e693e0..05c6b07fd 100644
--- a/src/ir/lubs.h
+++ b/src/ir/lubs.h
@@ -128,6 +128,19 @@ private:
   std::unordered_set<RefNull*> nulls;
 };
 
+namespace LUB {
+
+// Given a function, computes a LUB for its results. The caller can then decide
+// to apply a refined type if we found one.
+//
+// This modifies the called function even if it fails to find a refined type as
+// it does a refinalize in order to be able to compute the new types. We could
+// roll back that change, but it's not harmful and can help, so we keep it
+// regardless.
+LUBFinder getResultsLUB(Function* func, Module& wasm);
+
+} // namespace LUB
+
 } // namespace wasm
 
 #endif // wasm_ir_lubs_h
diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index 89e47708b..999607ae2 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -611,94 +611,22 @@ private:
   bool refineReturnTypes(Function* func,
                          const std::vector<Call*>& calls,
                          Module* module) {
-    if (!module->features.hasGC()) {
-      return false;
-    }
-
-    Type originalType = func->getResults();
-    if (!originalType.hasRef()) {
-      // Nothing to refine.
-      return false;
-    }
-
-    // Before we do anything, we must refinalize the function, because otherwise
-    // its body may contain a block with a forced type,
-    //
-    // (func (result X)
-    //  (block (result X)
-    //   (..content with more specific type Y..)
-    //  )
-    ReFinalize().walkFunctionInModule(func, module);
-
-    LUBFinder lub;
-    lub.noteUpdatableExpression(func->body);
-    if (lub.getBestPossible() == originalType) {
-      return false;
-    }
-
-    // Scan the body and look at the returns. First, return expressions.
-    for (auto* ret : FindAll<Return>(func->body).list) {
-      lub.noteUpdatableExpression(ret->value);
-      if (lub.getBestPossible() == originalType) {
-        return false;
-      }
-    }
-
-    // Process return_calls and call_refs. Unlike return expressions which we
-    // just handled, these only get a type to update, not a value.
-    auto processReturnType = [&](Type type) {
-      // Return whether we still look ok to do the optimization. If this is
-      // false then we can stop here.
-      lub.note(type);
-      return lub.getBestPossible() != originalType;
-    };
-
-    for (auto* call : FindAll<Call>(func->body).list) {
-      if (call->isReturn &&
-          !processReturnType(module->getFunction(call->target)->getResults())) {
-        return false;
-      }
-    }
-    for (auto* call : FindAll<CallIndirect>(func->body).list) {
-      if (call->isReturn &&
-          !processReturnType(call->heapType.getSignature().results)) {
-        return false;
-      }
-    }
-    for (auto* call : FindAll<CallRef>(func->body).list) {
-      if (call->isReturn) {
-        auto targetType = call->target->type;
-        if (targetType == Type::unreachable) {
-          continue;
-        }
-        if (!processReturnType(
-              targetType.getHeapType().getSignature().results)) {
-          return false;
-        }
-      }
-    }
-
-    // If the refined type is unreachable then nothing actually returns from
-    // this function.
-    // TODO: We can propagate that to the outside, and not just for GC.
+    auto lub = LUB::getResultsLUB(func, *module);
     if (!lub.noted()) {
       return false;
     }
-
     auto newType = lub.getBestPossible();
-    if (newType == originalType) {
-      return false;
-    }
-
-    // Success. Update the type, and the calls.
-    func->setResults(newType);
-    for (auto* call : calls) {
-      if (call->type != Type::unreachable) {
-        call->type = newType;
+    if (newType != func->getResults()) {
+      lub.updateNulls();
+      func->setResults(newType);
+      for (auto* call : calls) {
+        if (call->type != Type::unreachable) {
+          call->type = newType;
+        }
       }
+      return true;
     }
-    lub.updateNulls();
-    return true;
+    return false;
   }
 };