[Wasm GC] Optimize reference identity checks (#3814)

*  Note that ref.cast has a fallthrough value.
*  Optimize ref.eq on identical inputs.

3 files changed, 53 insertions, 0 deletions

diff --git a/src/ir/find_all.h b/src/ir/find_all.h
index bd3d265e2..77edafe89 100644
--- a/src/ir/find_all.h
+++ b/src/ir/find_all.h
@@ -40,6 +40,8 @@ template<typename T> struct FindAll {
     finder.list = &list;
     finder.walk(ast);
   }
+
+  bool has() { return !list.empty(); }
 };
 
 // Find all pointers to instances of a certain node type
diff --git a/src/ir/properties.h b/src/ir/properties.h
index c4deac9d5..379a85708 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -223,6 +223,9 @@ inline Index getZeroExtBits(Expression* curr) {
 // and other operations that receive a value and let it flow through them. If
 // there is no value falling through, returns the node itself (as that is the
 // value that trivially falls through, with 0 steps in the middle).
+// Note that this returns the value that would fall through if one does in fact
+// do so. For example, the final element in a block may not fall through if we
+// hit a return or a trap or an exception is thrown before we get there.
 inline Expression* getFallthrough(Expression* curr,
                                   const PassOptions& passOptions,
                                   FeatureSet features) {
@@ -259,6 +262,8 @@ inline Expression* getFallthrough(Expression* curr,
     if (!EffectAnalyzer(passOptions, features, tryy->body).throws) {
       return getFallthrough(tryy->body, passOptions, features);
     }
+  } else if (auto* as = curr->dynCast<RefCast>()) {
+    return getFallthrough(as->ref, passOptions, features);
   } else if (auto* as = curr->dynCast<RefAs>()) {
     return getFallthrough(as->value, passOptions, features);
   } else if (auto* br = curr->dynCast<BrOn>()) {
diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 1f55c0a4b..3c29f7042 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -26,6 +26,7 @@
 #include <ir/bits.h>
 #include <ir/cost.h>
 #include <ir/effects.h>
+#include <ir/find_all.h>
 #include <ir/gc-type-utils.h>
 #include <ir/literal-utils.h>
 #include <ir/load-utils.h>
@@ -992,6 +993,14 @@ struct OptimizeInstructions
     }
   }
 
+  void visitRefEq(RefEq* curr) {
+    // Identical references compare equal.
+    if (areConsecutiveInputsEqual(curr->left, curr->right)) {
+      replaceCurrent(
+        Builder(*getModule()).makeConst(Literal::makeOne(Type::i32)));
+    }
+  }
+
   // If an instruction traps on a null input, there is no need for a
   // ref.as_non_null on that input: we will trap either way (and the binaryen
   // optimizer does not differentiate traps).
@@ -1170,6 +1179,43 @@ private:
   // Information about our locals
   std::vector<LocalInfo> localInfo;
 
+  // Check if two consecutive inputs to an instruction are equal. As they are
+  // consecutive, no code can execeute in between them, which simplies the
+  // problem here (and which is the case we care about in this pass, which does
+  // simple peephole optimizations - all we care about is a single instruction
+  // at a time, and its inputs).
+  bool areConsecutiveInputsEqual(Expression* left, Expression* right) {
+    // First, check for side effects. If there are any, then we can't even
+    // assume things like local.get's of the same index being identical.
+    PassOptions passOptions = getPassOptions();
+    if (EffectAnalyzer(passOptions, getModule()->features, left)
+          .hasSideEffects() ||
+        EffectAnalyzer(passOptions, getModule()->features, right)
+          .hasSideEffects()) {
+      return false;
+    }
+    // Ignore extraneous things and compare them structurally.
+    left = Properties::getFallthrough(left, passOptions, getModule()->features);
+    right =
+      Properties::getFallthrough(right, passOptions, getModule()->features);
+    if (!ExpressionAnalyzer::equal(left, right)) {
+      return false;
+    }
+    // Reference equality also needs to check for allocation, which is *not* a
+    // side effect, but which results in different references:
+    // repeatedly calling (struct.new $foo)'s output will return different
+    // results (while i32.const etc. of course does not).
+    // Note that allocations may have appeared in the original inputs to this
+    // function, and skipped when we focused on what falls through; since there
+    // are no side effects, any allocations there cannot reach the fallthrough.
+    if (left->type.isRef()) {
+      if (FindAll<StructNew>(left).has() || FindAll<ArrayNew>(left).has()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   // Canonicalizing the order of a symmetric binary helps us
   // write more concise pattern matching code elsewhere.
   void canonicalize(Binary* binary) {