Update RemoveUnusedBrs for exact references

scripts/test/fuzzing.py

@@ -107,6 +107,7 @@
     'local-subtyping-exact.wast',
     'remove-unused-types-exact.wast',
     'coalesce-locals-exact.wast',
+    'remove-unused-brs-exact.wast',
 ]
 
 

src/ir/type-updating.cpp

@@ -205,9 +205,7 @@ void GlobalTypeRewriter::mapTypes(const TypeMap& oldToNewTypes) {
 
     Type getNew(Type type) {
       if (type.isRef()) {
-        return Type(getNew(type.getHeapType()),
-                    type.getNullability(),
-                    type.getExactness());
+        return type.with(getNew(type.getHeapType()));
       }
       if (type.isTuple()) {
         auto tuple = type.getTuple();
@@ -467,7 +465,7 @@ void handleNonDefaultableLocals(Function* func, Module& wasm) {
 Type getValidLocalType(Type type, FeatureSet features) {
   assert(type.isConcrete());
   if (type.isNonNullable()) {
-    return Type(type.getHeapType(), Nullable, type.getExactness());
+    return type.with(Nullable);
   }
   if (type.isTuple()) {
     std::vector<Type> elems(type.size());

src/passes/LocalSubtyping.cpp

@@ -152,8 +152,7 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
         // Remove non-nullability if we disallow that in locals.
         if (newType.isNonNullable()) {
           if (cannotBeNonNullable.count(i)) {
-            newType =
-              Type(newType.getHeapType(), Nullable, newType.getExactness());
+            newType = newType.with(Nullable);
           }
         } else if (!newType.isDefaultable()) {
           // Aside from the case we just handled of allowed non-nullability, we

src/passes/RemoveUnusedBrs.cpp

@@ -859,8 +859,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           if (Type::isSubType(expr->type, type)) {
             return expr;
           }
-          if (HeapType::isSubType(expr->type.getHeapType(),
-                                  type.getHeapType())) {
+          if (type.isNonNullable() && expr->type.isNullable() &&
+              Type::isSubType(expr->type.with(NonNullable), type)) {
             return builder.makeRefAs(RefAsNonNull, expr);
           }
           return builder.makeRefCast(expr, type);

src/wasm-type.h

@@ -398,6 +398,17 @@ class Type {
     return isExact() ? Exact : Inexact;
   }
 
+  // Return a new reference type with some part updated to the specified value.
+  Type with(HeapType heapType) {
+    return Type(heapType, getNullability(), getExactness());
+  }
+  Type with(Nullability nullability) {
+    return Type(getHeapType(), nullability, getExactness());
+  }
+  Type with(Exactness exactness) {
+    return Type(getHeapType(), getNullability(), exactness);
+  }
+
 private:
   template<bool (Type::*pred)() const> bool hasPredicate() {
     for (const auto& type : *this) {

src/wasm/wasm.cpp

@@ -1076,8 +1076,7 @@ void BrOn::finalize() {
   switch (op) {
     case BrOnNull:
       // If we do not branch, we flow out the existing value as non-null.
-      type =
-        Type(ref->type.getHeapType(), NonNullable, ref->type.getExactness());
+      type = ref->type.with(NonNullable);
       break;
     case BrOnNonNull:
       // If we do not branch, we flow out nothing (the spec could also have had
@@ -1087,8 +1086,7 @@ void BrOn::finalize() {
     case BrOnCast:
       if (castType.isNullable()) {
         // Nulls take the branch, so the result is non-nullable.
-        type =
-          Type(ref->type.getHeapType(), NonNullable, ref->type.getExactness());
+        type = ref->type.with(NonNullable);
       } else {
         // Nulls do not take the branch, so the result is non-nullable only if
         // the input is.
@@ -1099,9 +1097,7 @@ void BrOn::finalize() {
       if (castType.isNullable()) {
         // Nulls do not take the branch, so the result is non-nullable only if
         // the input is.
-        type = Type(castType.getHeapType(),
-                    ref->type.getNullability(),
-                    castType.getExactness());
+        type = castType.with(ref->type.getNullability());
       } else {
         // Nulls take the branch, so the result is non-nullable.
         type = castType;
@@ -1124,14 +1120,11 @@ Type BrOn::getSentType() {
         return Type::unreachable;
       }
       // BrOnNonNull sends the non-nullable type on the branch.
-      return Type(
-        ref->type.getHeapType(), NonNullable, ref->type.getExactness());
+      return ref->type.with(NonNullable);
     case BrOnCast:
       // The same as the result type of br_on_cast_fail.
       if (castType.isNullable()) {
-        return Type(castType.getHeapType(),
-                    ref->type.getNullability(),
-                    castType.getExactness());
+        return castType.with(ref->type.getNullability());
       } else {
         return castType;
       }
@@ -1141,8 +1134,7 @@ Type BrOn::getSentType() {
         return Type::unreachable;
       }
       if (castType.isNullable()) {
-        return Type(
-          ref->type.getHeapType(), NonNullable, ref->type.getExactness());
+        return ref->type.with(NonNullable);
       } else {
         return ref->type;
       }
@@ -1317,7 +1309,7 @@ void RefAs::finalize() {
   auto valHeapType = value->type.getHeapType();
   switch (op) {
     case RefAsNonNull:
-      type = Type(valHeapType, NonNullable, value->type.getExactness());
+      type = value->type.with(NonNullable);
       break;
     case AnyConvertExtern:
       type = Type(HeapTypes::any.getBasic(valHeapType.getShared()),

test/lit/passes/remove-unused-brs-exact.wast

@@ -0,0 +1,40 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+
+;; RUN: wasm-opt %s -all --remove-unused-brs -S -o - | filecheck %s
+
+;; Check that we optimize the cast correctly when the fallthrough has exact
+;; type. In particular, we should not insert a ref.as_non_null, which would
+;; trap.
+
+(module
+ ;; CHECK:      (func $br_on_cast_fail (type $0) (param $0 (exact nullref))
+ ;; CHECK-NEXT:  (local $1 nullref)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (block $block
+ ;; CHECK-NEXT:    (drop
+ ;; CHECK-NEXT:     (ref.cast (exact nullref)
+ ;; CHECK-NEXT:      (local.tee $1
+ ;; CHECK-NEXT:       (local.get $0)
+ ;; CHECK-NEXT:      )
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:    (return)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $br_on_cast_fail (param (ref null exact none))
+  (local $1 nullref)
+  (drop
+   (block $block (result (ref none))
+    (drop
+     (br_on_cast_fail $block nullref nullref
+      (local.tee $1
+       (local.get 0)
+      )
+     )
+    )
+    (return)
+   )
+  )
+ )
+)