Fix LocalCSE's usable local selection (#2638)

Now that we have subtypes, we cannot reuse any local that contains the
same expression, because that local's type can be a supertype. For
example:
```
(local $0 anyref)
(local $1 nullref)
...
(local.set $0 (ref.null))
(local.set $1 (ref.null)) ;; cannot be replaced with (local.get $0)
```

This extends `usables` map's key to contain both `HashedExpression` and
the local's type, so we can get the right usable local in presence of
subtypes.

4 files changed, 67 insertions, 27 deletions

diff --git a/src/ir/hashed.h b/src/ir/hashed.h
index fe0a0b958..3068966c5 100644
--- a/src/ir/hashed.h
+++ b/src/ir/hashed.h
@@ -38,25 +38,6 @@ struct HashedExpression {
     : expr(other.expr), hash(other.hash) {}
 };
 
-struct ExpressionHasher {
-  HashType operator()(const HashedExpression value) const { return value.hash; }
-};
-
-struct ExpressionComparer {
-  bool operator()(const HashedExpression a, const HashedExpression b) const {
-    if (a.hash != b.hash) {
-      return false;
-    }
-    return ExpressionAnalyzer::equal(a.expr, b.expr);
-  }
-};
-
-template<typename T>
-class HashedExpressionMap
-  : public std::
-      unordered_map<HashedExpression, T, ExpressionHasher, ExpressionComparer> {
-};
-
 // A pass that hashes all functions
 
 struct FunctionHasher : public WalkerPass<PostWalker<FunctionHasher>> {
diff --git a/src/passes/LocalCSE.cpp b/src/passes/LocalCSE.cpp
index f7c1cda26..d6717a1ed 100644
--- a/src/passes/LocalCSE.cpp
+++ b/src/passes/LocalCSE.cpp
@@ -54,6 +54,28 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
 
   Pass* create() override { return new LocalCSE(); }
 
+  struct Usable {
+    HashedExpression hashed;
+    Type localType;
+    Usable(HashedExpression hashed, Type localType)
+      : hashed(hashed), localType(localType) {}
+  };
+
+  struct UsableHasher {
+    HashType operator()(const Usable value) const {
+      return rehash(value.hashed.hash, value.localType.getID());
+    }
+  };
+
+  struct UsableComparer {
+    bool operator()(const Usable a, const Usable b) const {
+      if (a.hashed.hash != b.hashed.hash || a.localType != b.localType) {
+        return false;
+      }
+      return ExpressionAnalyzer::equal(a.hashed.expr, b.hashed.expr);
+    }
+  };
+
   // information for an expression we can reuse
   struct UsableInfo {
     Expression* value; // the value we can reuse
@@ -68,7 +90,9 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
   };
 
   // a list of usables in a linear execution trace
-  typedef HashedExpressionMap<UsableInfo> Usables;
+  class Usables
+    : public std::
+        unordered_map<Usable, UsableInfo, UsableHasher, UsableComparer> {};
 
   // locals in current linear execution trace, which we try to sink
   Usables usables;
@@ -104,7 +128,7 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
   // taken into account.
   void checkInvalidations(EffectAnalyzer& effects, Expression* curr = nullptr) {
     // TODO: this is O(bad)
-    std::vector<HashedExpression> invalidated;
+    std::vector<Usable> invalidated;
     for (auto& sinkable : usables) {
       // Check invalidations of the values we may want to use.
       if (effects.invalidates(sinkable.second.effects)) {
@@ -185,8 +209,8 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
       // consider the value
       auto* value = set->value;
       if (isRelevant(value)) {
-        HashedExpression hashed(value);
-        auto iter = usables.find(hashed);
+        Usable usable(value, func->getLocalType(set->index));
+        auto iter = usables.find(usable);
         if (iter != usables.end()) {
           // already exists in the table, this is good to reuse
           auto& info = iter->second;
@@ -197,7 +221,7 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
         } else {
           // not in table, add this, maybe we can help others later
           usables.emplace(std::make_pair(
-            hashed,
+            usable,
             UsableInfo(
               value, set->index, getPassOptions(), getModule()->features)));
         }
diff --git a/test/passes/flatten_local-cse_all-features.txt b/test/passes/flatten_local-cse_all-features.txt
index 0697e27a5..fc2dff3aa 100644
--- a/test/passes/flatten_local-cse_all-features.txt
+++ b/test/passes/flatten_local-cse_all-features.txt
@@ -762,6 +762,7 @@
  )
 )
 (module
+ (type $none_=>_none (func))
  (type $none_=>_funcref (func (result funcref)))
  (func $subtype-test (; 0 ;) (result funcref)
   (local $0 nullref)
@@ -789,4 +790,23 @@
    (local.get $2)
   )
  )
+ (func $test (; 1 ;)
+  (local $0 anyref)
+  (local $1 nullref)
+  (local $2 nullref)
+  (block $label$1
+   (local.set $0
+    (ref.null)
+   )
+   (local.set $1
+    (ref.null)
+   )
+  )
+  (local.set $2
+   (local.get $1)
+  )
+  (drop
+   (local.get $1)
+  )
+ )
 )
diff --git a/test/passes/flatten_local-cse_all-features.wast b/test/passes/flatten_local-cse_all-features.wast
index 6ca154278..112dfe647 100644
--- a/test/passes/flatten_local-cse_all-features.wast
+++ b/test/passes/flatten_local-cse_all-features.wast
@@ -288,14 +288,29 @@
  )
 )
 
-;; After --flatten, there will be a series of chain copies between multiple
-;; locals, but some of the locals will be nullref type and others funcref type.
-;; We cannot make locals of different types a common subexpression.
 (module
+ ;; After --flatten, there will be a series of chain copies between multiple
+ ;; locals, but some of the locals will be nullref type and others funcref type.
+ ;; We cannot make locals of different types a common subexpression.
  (func $subtype-test (result funcref)
   (nop)
   (loop $label$1 (result nullref)
    (ref.null)
   )
  )
+
+ (func $test
+  (local $0 anyref)
+  (drop
+   (block $label$1 (result nullref)
+    (local.set $0
+     (ref.null)
+    )
+    ;; After --flatten, this will be assigned to a local of nullref type. After
+    ;; --local-cse, even if we set (ref.null) to local $0 above, this should not
+    ;; be replaced with $0, because it is of type anyref.
+    (ref.null)
+   )
+  )
+ )
 )