Fix reuse of constant nodes in Precompute (#2764)

Previously we tried to reuse `Const` node if a precomputed value is a
constant node. But now we have two more kinds of constant node
(`RefNull` and `RefFunc`), so we shouldn't reuse them interchangeably,
meaning we shouldn't try to reuse a `Const` node when the value at hand
is a `RefNull`. This correctly checks the type of node and tries to
reuse only if the types of nodes match.

Fixes #2759.

3 files changed, 156 insertions, 28 deletions

diff --git a/src/passes/Precompute.cpp b/src/passes/Precompute.cpp
index 0a4a691f2..d809b8b43 100644
--- a/src/passes/Precompute.cpp
+++ b/src/passes/Precompute.cpp
@@ -174,6 +174,37 @@ struct Precompute
     } while (propagate && worked);
   }
 
+  template<typename T> void reuseConstantNode(T* curr, Flow flow) {
+    if (flow.values.isConcrete()) {
+      // reuse a const / ref.null / ref.func node if there is one
+      if (curr->value && flow.values.size() == 1) {
+        Literal singleValue = flow.getSingleValue();
+        if (singleValue.type.isNumber()) {
+          if (auto* c = curr->value->template dynCast<Const>()) {
+            c->value = singleValue;
+            c->finalize();
+            curr->finalize();
+            return;
+          }
+        } else if (singleValue.type == Type::nullref &&
+                   curr->value->template is<RefNull>()) {
+          return;
+        } else if (singleValue.type == Type::funcref) {
+          if (auto* r = curr->value->template dynCast<RefFunc>()) {
+            r->func = singleValue.getFunc();
+            r->finalize();
+            curr->finalize();
+            return;
+          }
+        }
+      }
+      curr->value = flow.getConstExpression(*getModule());
+    } else {
+      curr->value = nullptr;
+    }
+    curr->finalize();
+  }
+
   void visitExpression(Expression* curr) {
     // TODO: if local.get, only replace with a constant if we don't care about
     // size...?
@@ -199,19 +230,7 @@ struct Precompute
         // this expression causes a return. if it's already a return, reuse the
         // node
         if (auto* ret = curr->dynCast<Return>()) {
-          if (flow.values.isConcrete()) {
-            // reuse a const value if there is one
-            if (ret->value && flow.values.size() == 1) {
-              if (auto* value = ret->value->dynCast<Const>()) {
-                value->value = flow.getSingleValue();
-                value->finalize();
-                return;
-              }
-            }
-            ret->value = flow.getConstExpression(*getModule());
-          } else {
-            ret->value = nullptr;
-          }
+          reuseConstantNode(ret, flow);
         } else {
           Builder builder(*getModule());
           replaceCurrent(builder.makeReturn(
@@ -225,21 +244,7 @@ struct Precompute
       if (auto* br = curr->dynCast<Break>()) {
         br->name = flow.breakTo;
         br->condition = nullptr;
-        if (flow.values.isConcrete()) {
-          // reuse a const value if there is one
-          if (br->value && flow.values.size() == 1) {
-            if (auto* value = br->value->dynCast<Const>()) {
-              value->value = flow.getSingleValue();
-              value->finalize();
-              br->finalize();
-              return;
-            }
-          }
-          br->value = flow.getConstExpression(*getModule());
-        } else {
-          br->value = nullptr;
-        }
-        br->finalize();
+        reuseConstantNode(br, flow);
       } else {
         Builder builder(*getModule());
         replaceCurrent(builder.makeBreak(
diff --git a/test/passes/precompute_all-features.txt b/test/passes/precompute_all-features.txt
index 9a6fea94b..34c9554dd 100644
--- a/test/passes/precompute_all-features.txt
+++ b/test/passes/precompute_all-features.txt
@@ -261,4 +261,54 @@
  (func $reftype-test (result nullref)
   (ref.null)
  )
+ (func $dummy
+  (nop)
+ )
+ (func $br_reuse_node
+  (drop
+   (block $l0 (result f32)
+    (drop
+     (block $l1
+      (global.set $global-mut
+       (i32.const 1)
+      )
+      (br $l0
+       (f32.const 3.5)
+      )
+     )
+    )
+    (f32.const 0)
+   )
+  )
+  (drop
+   (block $l2 (result nullref)
+    (drop
+     (block $l3
+      (global.set $global-mut
+       (i32.const 1)
+      )
+      (br $l2
+       (ref.null)
+      )
+     )
+    )
+    (ref.null)
+   )
+  )
+  (drop
+   (block $l4 (result funcref)
+    (drop
+     (block $l5
+      (global.set $global-mut
+       (i32.const 1)
+      )
+      (br $l4
+       (ref.func $dummy)
+      )
+     )
+    )
+    (ref.null)
+   )
+  )
+ )
 )
diff --git a/test/passes/precompute_all-features.wast b/test/passes/precompute_all-features.wast
index da755bde6..2a59ea290 100644
--- a/test/passes/precompute_all-features.wast
+++ b/test/passes/precompute_all-features.wast
@@ -381,4 +381,77 @@
   (func $reftype-test (result nullref)
     (ref.null)
   )
+
+  ;; Check if constant nodes (const, ref.null, and ref.func) are correctly
+  ;; reused. (We shouldn't reuse a const node for something like ref.null, which
+  ;; will incorrectly cause an expression like 'nullref.const'.)
+  (func $dummy)
+  (func $br_reuse_node
+   (drop
+    (block $l0 (result f32)
+     (drop
+      (block $l1 (result i32)
+       (global.set $global-mut
+        (i32.const 1)
+       )
+       (br_if $l1
+        (i32.const 1)
+        (f32.lt
+         (br_if $l0
+          (f32.const 3.5)
+          (i32.const 1)
+         )
+         (f32.const 3)
+        )
+       )
+      )
+     )
+     (f32.const 0)
+    )
+   )
+
+   (drop
+    (block $l2 (result nullref)
+     (drop
+      (block $l3 (result i32)
+       (global.set $global-mut
+        (i32.const 1)
+       )
+       (br_if $l3
+        (i32.const 1)
+        (ref.is_null
+         (br_if $l2
+          (ref.null)
+          (i32.const 3)
+         )
+        )
+       )
+      )
+     )
+     (ref.null)
+    )
+   )
+
+   (drop
+    (block $l4 (result funcref)
+     (drop
+      (block $l5 (result i32)
+       (global.set $global-mut
+        (i32.const 1)
+       )
+       (br_if $l5
+        (i32.const 1)
+        (ref.is_null
+         (br_if $l4
+          (ref.func $dummy)
+          (i32.const 3)
+         )
+        )
+       )
+      )
+     )
+     (ref.null)
+    )
+   )
+  )
 )