Fix br_on_non_null value

Author: osa1

src/passes/Flatten.cpp

@@ -37,6 +37,7 @@
 // local, but in principle there's no reason it couldn't be. For now, error on
 // this.
 
+#include "wasm-type.h"
 #include <ir/branch-utils.h>
 #include <ir/effects.h>
 #include <ir/eh-utils.h>
@@ -335,10 +336,20 @@ struct Flatten
               builder.makeUnary(UnaryOp::EqZInt32,
                                 builder.makeLocalGet(isNullTemp, Type::i32))));
 
-            replaceCurrent(builder.makeBreak(
-              br->name,
-              nullptr,
-              builder.makeLocalGet(isNotNullTemp, Type::i32)));
+            Index breakTargetTemp = getTempForBreakTarget(
+              br->name, nullableType.with(Nullability::NonNullable));
+
+            std::vector<Expression*> successBlock;
+            successBlock.push_back(builder.makeLocalSet(
+              breakTargetTemp,
+              builder.makeRefAs(
+                RefAsOp::RefAsNonNull,
+                builder.makeLocalGet(nullableTemp, nullableType))));
+            successBlock.push_back(builder.makeBreak(br->name));
+
+            replaceCurrent(
+              builder.makeIf(builder.makeLocalGet(isNotNullTemp, Type::i32),
+                             builder.makeBlock(successBlock)));
           }
         } else if (br->op == BrOnCast || br->op == BrOnCastFail) {
           auto sourceType = br->ref->type;

test/lit/passes/flatten_br_on_non_null.wast

@@ -15,8 +15,8 @@
   ;; CHECK-NEXT:  (local $2 (ref null $s))
   ;; CHECK-NEXT:  (local $3 i32)
   ;; CHECK-NEXT:  (local $4 i32)
-  ;; CHECK-NEXT:  (local $5 (ref (exact $s)))
-  ;; CHECK-NEXT:  (local $6 (ref null $s))
+  ;; CHECK-NEXT:  (local $5 (ref null $s))
+  ;; CHECK-NEXT:  (local $6 (ref (exact $s)))
   ;; CHECK-NEXT:  (local $7 (ref $s))
   ;; CHECK-NEXT:  (block $label0
   ;; CHECK-NEXT:   (local.set $1
@@ -35,19 +35,27 @@
   ;; CHECK-NEXT:     (local.get $3)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (br_if $label0
+  ;; CHECK-NEXT:   (if
   ;; CHECK-NEXT:    (local.get $4)
+  ;; CHECK-NEXT:    (then
+  ;; CHECK-NEXT:     (local.set $5
+  ;; CHECK-NEXT:      (ref.as_non_null
+  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (br $label0)
+  ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $5
+  ;; CHECK-NEXT:   (local.set $6
   ;; CHECK-NEXT:    (struct.new_default $s)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $6
-  ;; CHECK-NEXT:    (local.get $5)
+  ;; CHECK-NEXT:   (local.set $5
+  ;; CHECK-NEXT:    (local.get $6)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (local.set $7
   ;; CHECK-NEXT:   (ref.as_non_null
-  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:    (local.get $5)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (return