From f8ec1128c49a1d30412377d6018f64c0dd474eac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Wed, 1 Oct 2025 12:05:21 +0100
Subject: [PATCH 01/11] Implement flattening for br_on_null

---
 src/passes/Flatten.cpp | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 1c2cfbcd536..46e2464c127 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -308,6 +308,37 @@ struct Flatten
           }
         }
 
+      } else if (auto* br = curr->dynCast<BrOn>()) {
+        if (br->op == BrOnOp::BrOnNull) {
+          auto nullableType = br->ref->type;
+          auto nonNullableType = nullableType.with(Nullability::NonNullable);
+          Index nonNullableTemp =
+            builder.addVar(getFunction(), nonNullableType);
+
+          Index nullableTemp = builder.addVar(getFunction(), nullableType);
+          ourPreludes.push_back(builder.makeLocalSet(nullableTemp, br->ref));
+
+          Index isNullTemp = builder.addVar(getFunction(), Type::i32);
+          ourPreludes.push_back(
+            builder.makeLocalSet(isNullTemp,
+                                 builder.makeRefIsNull(builder.makeLocalGet(
+                                   nullableTemp, nullableType))));
+
+          ourPreludes.push_back(builder.makeBreak(
+            br->name, nullptr, builder.makeLocalGet(isNullTemp, Type::i32)));
+
+          ourPreludes.push_back(builder.makeLocalSet(
+            nonNullableTemp,
+            builder.makeRefAs(
+              RefAsOp::RefAsNonNull,
+              builder.makeLocalGet(nullableTemp, nullableType))));
+
+          replaceCurrent(
+            builder.makeLocalGet(nonNullableTemp, nonNullableType));
+        } else {
+          Fatal() << "Unsupported instruction for Flatten: "
+                  << getExpressionName(curr);
+        }
       } else if (auto* sw = curr->dynCast<Switch>()) {
         if (sw->value) {
           auto type = sw->value->type;
@@ -333,7 +364,7 @@ struct Flatten
       }
     }
 
-    if (curr->is<BrOn>() || curr->is<TryTable>()) {
+    if (curr->is<TryTable>()) {
       Fatal() << "Unsupported instruction for Flatten: "
               << getExpressionName(curr);
     }

From 59835796f0ea7c629c1832fff64dd4b9745e2375 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Wed, 1 Oct 2025 12:13:31 +0100
Subject: [PATCH 02/11] Implement flattening for br_on_non_null

---
 src/passes/Flatten.cpp | 37 +++++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 46e2464c127..21df1fb41b5 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -309,7 +309,7 @@ struct Flatten
         }
 
       } else if (auto* br = curr->dynCast<BrOn>()) {
-        if (br->op == BrOnOp::BrOnNull) {
+        if (br->op == BrOnOp::BrOnNull || br->op == BrOnOp::BrOnNonNull) {
           auto nullableType = br->ref->type;
           auto nonNullableType = nullableType.with(Nullability::NonNullable);
           Index nonNullableTemp =
@@ -324,17 +324,30 @@ struct Flatten
                                  builder.makeRefIsNull(builder.makeLocalGet(
                                    nullableTemp, nullableType))));
 
-          ourPreludes.push_back(builder.makeBreak(
-            br->name, nullptr, builder.makeLocalGet(isNullTemp, Type::i32)));
-
-          ourPreludes.push_back(builder.makeLocalSet(
-            nonNullableTemp,
-            builder.makeRefAs(
-              RefAsOp::RefAsNonNull,
-              builder.makeLocalGet(nullableTemp, nullableType))));
-
-          replaceCurrent(
-            builder.makeLocalGet(nonNullableTemp, nonNullableType));
+          if (br->op == BrOnOp::BrOnNull) {
+            ourPreludes.push_back(builder.makeBreak(
+              br->name, nullptr, builder.makeLocalGet(isNullTemp, Type::i32)));
+
+            ourPreludes.push_back(builder.makeLocalSet(
+              nonNullableTemp,
+              builder.makeRefAs(
+                RefAsOp::RefAsNonNull,
+                builder.makeLocalGet(nullableTemp, nullableType))));
+
+            replaceCurrent(
+              builder.makeLocalGet(nonNullableTemp, nonNullableType));
+          } else { // BrOnNonNull
+            Index isNotNullTemp = builder.addVar(getFunction(), Type::i32);
+            ourPreludes.push_back(builder.makeLocalSet(
+              isNotNullTemp,
+              builder.makeUnary(UnaryOp::EqZInt32,
+                                builder.makeLocalGet(isNullTemp, Type::i32))));
+
+            replaceCurrent(builder.makeBreak(
+              br->name,
+              nullptr,
+              builder.makeLocalGet(isNotNullTemp, Type::i32)));
+          }
         } else {
           Fatal() << "Unsupported instruction for Flatten: "
                   << getExpressionName(curr);

From a0a5babdcb07dd915729adcfaafc8b257bfc5d5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Wed, 1 Oct 2025 13:58:34 +0100
Subject: [PATCH 03/11] Implement br_on_cast

---
 src/passes/Flatten.cpp | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 21df1fb41b5..2d4a12e12ac 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -348,9 +348,45 @@ struct Flatten
               nullptr,
               builder.makeLocalGet(isNotNullTemp, Type::i32)));
           }
+        } else if (br->op == BrOnCast) {
+          auto sourceType = br->ref->type;
+          auto targetType = br->castType;
+
+          Index sourceTypeTemp = builder.addVar(getFunction(), sourceType);
+          ourPreludes.push_back(builder.makeLocalSet(sourceTypeTemp, br->ref));
+
+          Index typeTestTemp = builder.addVar(getFunction(), Type::i32);
+          ourPreludes.push_back(builder.makeLocalSet(
+            typeTestTemp,
+            builder.makeRefTest(
+              builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
+
+          Expression* failValue =
+            builder.makeLocalGet(sourceTypeTemp, sourceType);
+          if (br->castType.isNullable()) {
+            failValue = builder.makeRefAs(RefAsOp::RefAsNonNull, failValue);
+          }
+
+          Index failTemp = builder.addVar(getFunction(), failValue->type);
+
+          Index targetTemp = getTempForBreakTarget(br->name, targetType);
+
+          std::vector<Expression*> successBlock;
+          successBlock.push_back(builder.makeLocalSet(
+            targetTemp,
+            builder.makeRefCast(
+              builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
+          successBlock.push_back(builder.makeBreak(br->name));
+
+          ourPreludes.push_back(
+            builder.makeIf(builder.makeLocalGet(typeTestTemp, Type::i32),
+                           builder.makeBlock(successBlock)));
+
+          ourPreludes.push_back(builder.makeLocalSet(failTemp, failValue));
+
+          replaceCurrent(builder.makeLocalGet(failTemp, failValue->type));
         } else {
-          Fatal() << "Unsupported instruction for Flatten: "
-                  << getExpressionName(curr);
+          Fatal() << "Unsupported instruction for Flatten: BrOn " << br->op;
         }
       } else if (auto* sw = curr->dynCast<Switch>()) {
         if (sw->value) {

From e12bf86ed20bbabf31501d58f87eb0fe8c41e0cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Wed, 1 Oct 2025 14:23:59 +0100
Subject: [PATCH 04/11] WIP

---
 src/passes/Flatten.cpp | 57 ++++++++++++++++++++++++++++++------------
 test.wast              | 24 ++++++++++++++++++
 2 files changed, 65 insertions(+), 16 deletions(-)
 create mode 100644 test.wast

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 2d4a12e12ac..797e4eb4745 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -336,7 +336,7 @@ struct Flatten
 
             replaceCurrent(
               builder.makeLocalGet(nonNullableTemp, nonNullableType));
-          } else { // BrOnNonNull
+          } else { // br_on_non_null
             Index isNotNullTemp = builder.addVar(getFunction(), Type::i32);
             ourPreludes.push_back(builder.makeLocalSet(
               isNotNullTemp,
@@ -348,7 +348,7 @@ struct Flatten
               nullptr,
               builder.makeLocalGet(isNotNullTemp, Type::i32)));
           }
-        } else if (br->op == BrOnCast) {
+        } else if (br->op == BrOnCast || br->op == BrOnCastFail) {
           auto sourceType = br->ref->type;
           auto targetType = br->castType;
 
@@ -367,24 +367,49 @@ struct Flatten
             failValue = builder.makeRefAs(RefAsOp::RefAsNonNull, failValue);
           }
 
-          Index failTemp = builder.addVar(getFunction(), failValue->type);
-
-          Index targetTemp = getTempForBreakTarget(br->name, targetType);
+          Index breakTargetTemp = getTempForBreakTarget(br->name, targetType);
+
+          if (br->op == BrOnCast) {
+            std::vector<Expression*> successBlock;
+            successBlock.push_back(builder.makeLocalSet(
+              breakTargetTemp,
+              builder.makeRefCast(
+                builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
+            successBlock.push_back(builder.makeBreak(br->name));
+
+            ourPreludes.push_back(
+              builder.makeIf(builder.makeLocalGet(typeTestTemp, Type::i32),
+                             builder.makeBlock(successBlock)));
+
+            Index failTemp = builder.addVar(getFunction(), failValue->type);
+            ourPreludes.push_back(builder.makeLocalSet(failTemp, failValue));
+            replaceCurrent(builder.makeLocalGet(failTemp, failValue->type));
+          } else { // br_on_cast_fail
+            std::vector<Expression*> failureBlock;
+            failureBlock.push_back(
+              builder.makeLocalSet(breakTargetTemp, failValue));
+            failureBlock.push_back(builder.makeBreak(br->name));
+
+            Index typeTestFailTemp = builder.addVar(getFunction(), Type::i32);
+            ourPreludes.push_back(builder.makeLocalSet(
+              typeTestFailTemp,
+              builder.makeUnary(
+                UnaryOp::EqZInt32,
+                builder.makeLocalGet(typeTestTemp, Type::i32))));
 
-          std::vector<Expression*> successBlock;
-          successBlock.push_back(builder.makeLocalSet(
-            targetTemp,
-            builder.makeRefCast(
-              builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
-          successBlock.push_back(builder.makeBreak(br->name));
+            ourPreludes.push_back(
+              builder.makeIf(builder.makeLocalGet(typeTestFailTemp, Type::i32),
+                             builder.makeBlock(failureBlock)));
 
-          ourPreludes.push_back(
-            builder.makeIf(builder.makeLocalGet(typeTestTemp, Type::i32),
-                           builder.makeBlock(successBlock)));
+            Index targetTypeTemp = builder.addVar(getFunction(), targetType);
+            ourPreludes.push_back(builder.makeLocalSet(
+              targetTypeTemp,
+              builder.makeRefCast(
+                builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
 
-          ourPreludes.push_back(builder.makeLocalSet(failTemp, failValue));
+            replaceCurrent(builder.makeLocalGet(targetTypeTemp, targetType));
+          }
 
-          replaceCurrent(builder.makeLocalGet(failTemp, failValue->type));
         } else {
           Fatal() << "Unsupported instruction for Flatten: BrOn " << br->op;
         }
diff --git a/test.wast b/test.wast
new file mode 100644
index 00000000000..260cfcec323
--- /dev/null
+++ b/test.wast
@@ -0,0 +1,24 @@
+(module
+ (type $s (sub (struct)))
+ (type $t (sub $s (struct)))
+
+  (func $br_on_null (param $x (ref null $s)) (result (ref $s))
+    (block $label0
+      (return (br_on_null $label0 (local.get $x))))
+    (struct.new $s))
+
+  (func $br_on_non_null (param $x (ref null $s)) (result (ref $s))
+    (return (block $label0 (result (ref $s))
+      (br_on_non_null $label0 (local.get $x))
+      (struct.new $s))))
+
+  (func $br_on_cast (param $x (ref $s)) (result (ref $t))
+    (return (block $label0 (result (ref $t))
+      (drop (br_on_cast $label0 (ref $s) (ref $t) (local.get $x)))
+      (struct.new $t))))
+
+  (func $br_on_cast_fail (param $x (ref $s)) (result (ref $t))
+    (drop (block $label0 (result (ref $s))
+      (return (br_on_cast_fail $label0 (ref $s) (ref $t) (local.get $x)))))
+    (struct.new $t))
+  )

From 1e2420d982ef0caf6fb9df7a8f2fe5355840ec3c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 09:33:38 +0100
Subject: [PATCH 05/11] Fix bug

---
 src/passes/Flatten.cpp | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index 797e4eb4745..e47e8c05cfc 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -361,15 +361,17 @@ struct Flatten
             builder.makeRefTest(
               builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));
 
+          // On cast failure the source type is made non-nullable if the target
+          // type is nullable.
           Expression* failValue =
             builder.makeLocalGet(sourceTypeTemp, sourceType);
           if (br->castType.isNullable()) {
             failValue = builder.makeRefAs(RefAsOp::RefAsNonNull, failValue);
           }
 
-          Index breakTargetTemp = getTempForBreakTarget(br->name, targetType);
-
           if (br->op == BrOnCast) {
+            Index breakTargetTemp = getTempForBreakTarget(br->name, targetType);
+
             std::vector<Expression*> successBlock;
             successBlock.push_back(builder.makeLocalSet(
               breakTargetTemp,
@@ -385,6 +387,9 @@ struct Flatten
             ourPreludes.push_back(builder.makeLocalSet(failTemp, failValue));
             replaceCurrent(builder.makeLocalGet(failTemp, failValue->type));
           } else { // br_on_cast_fail
+            Index breakTargetTemp =
+              getTempForBreakTarget(br->name, failValue->type);
+
             std::vector<Expression*> failureBlock;
             failureBlock.push_back(
               builder.makeLocalSet(breakTargetTemp, failValue));

From cd31b32ca91ddf8b4b7c3b70313bd3c3803052ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 09:41:04 +0100
Subject: [PATCH 06/11] Add tests

---
 test/lit/passes/flatten_br_on_cast.wast      | 77 ++++++++++++++++
 test/lit/passes/flatten_br_on_cast_fail.wast | 92 ++++++++++++++++++++
 test/lit/passes/flatten_br_on_non_null.wast  | 64 ++++++++++++++
 test/lit/passes/flatten_br_on_null.wast      | 69 +++++++++++++++
 4 files changed, 302 insertions(+)
 create mode 100644 test/lit/passes/flatten_br_on_cast.wast
 create mode 100644 test/lit/passes/flatten_br_on_cast_fail.wast
 create mode 100644 test/lit/passes/flatten_br_on_non_null.wast
 create mode 100644 test/lit/passes/flatten_br_on_null.wast

diff --git a/test/lit/passes/flatten_br_on_cast.wast b/test/lit/passes/flatten_br_on_cast.wast
new file mode 100644
index 00000000000..dc267584550
--- /dev/null
+++ b/test/lit/passes/flatten_br_on_cast.wast
@@ -0,0 +1,77 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_passes_tests_to_lit.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $s (sub (struct)))
+ (type $s (sub (struct)))
+ ;; CHECK:      (type $t (sub $s (struct)))
+ (type $t (sub $s (struct)))
+
+  ;; CHECK:      (type $2 (func (param (ref $s)) (result (ref $t))))
+
+  ;; CHECK:      (func $br_on_cast (type $2) (param $x (ref $s)) (result (ref $t))
+  ;; CHECK-NEXT:  (local $1 (ref $s))
+  ;; CHECK-NEXT:  (local $2 (ref $s))
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 (ref null $t))
+  ;; CHECK-NEXT:  (local $5 (ref $s))
+  ;; CHECK-NEXT:  (local $6 (ref $s))
+  ;; CHECK-NEXT:  (local $7 (ref (exact $t)))
+  ;; CHECK-NEXT:  (local $8 (ref $t))
+  ;; CHECK-NEXT:  (block $label0
+  ;; CHECK-NEXT:   (local.set $1
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $2
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $3
+  ;; CHECK-NEXT:    (ref.test (ref $t)
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (if
+  ;; CHECK-NEXT:    (local.get $3)
+  ;; CHECK-NEXT:    (then
+  ;; CHECK-NEXT:     (local.set $4
+  ;; CHECK-NEXT:      (ref.cast (ref $t)
+  ;; 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.get $2)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:    (local.get $5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $7
+  ;; CHECK-NEXT:    (struct.new_default $t)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $4
+  ;; CHECK-NEXT:    (local.get $7)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $8
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $4)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return
+  ;; CHECK-NEXT:   (local.get $8)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $br_on_cast (param $x (ref $s)) (result (ref $t))
+    (return (block $label0 (result (ref $t))
+      (drop (br_on_cast $label0 (ref $s) (ref $t) (local.get $x)))
+      (struct.new $t))))
+)
+
diff --git a/test/lit/passes/flatten_br_on_cast_fail.wast b/test/lit/passes/flatten_br_on_cast_fail.wast
new file mode 100644
index 00000000000..2c55763a615
--- /dev/null
+++ b/test/lit/passes/flatten_br_on_cast_fail.wast
@@ -0,0 +1,92 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_passes_tests_to_lit.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $s (sub (struct)))
+ (type $s (sub (struct)))
+ ;; CHECK:      (type $t (sub $s (struct)))
+ (type $t (sub $s (struct)))
+
+  ;; CHECK:      (type $2 (func (param (ref $s)) (result (ref $t))))
+
+  ;; CHECK:      (func $br_on_cast_fail (type $2) (param $x (ref $s)) (result (ref $t))
+  ;; CHECK-NEXT:  (local $1 (ref $s))
+  ;; CHECK-NEXT:  (local $2 (ref $s))
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 (ref null $s))
+  ;; CHECK-NEXT:  (local $5 i32)
+  ;; CHECK-NEXT:  (local $6 (ref $t))
+  ;; CHECK-NEXT:  (local $7 (ref $t))
+  ;; CHECK-NEXT:  (local $8 (ref $s))
+  ;; CHECK-NEXT:  (local $9 (ref (exact $t)))
+  ;; CHECK-NEXT:  (local $10 (ref $t))
+  ;; CHECK-NEXT:  (local $11 (ref $t))
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (block $label0
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (local.get $1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (ref.test (ref $t)
+  ;; CHECK-NEXT:      (local.get $2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $5
+  ;; CHECK-NEXT:     (i32.eqz
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (local.get $5)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (local.set $4
+  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (br $label0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $6
+  ;; CHECK-NEXT:     (ref.cast (ref $t)
+  ;; CHECK-NEXT:      (local.get $2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $7
+  ;; CHECK-NEXT:     (local.get $6)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (local.get $7)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $8
+  ;; CHECK-NEXT:    (ref.as_non_null
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (local.get $8)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $9
+  ;; CHECK-NEXT:    (struct.new_default $t)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $10
+  ;; CHECK-NEXT:    (local.get $9)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $11
+  ;; CHECK-NEXT:   (local.get $10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return
+  ;; CHECK-NEXT:   (local.get $11)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $br_on_cast_fail (param $x (ref $s)) (result (ref $t))
+    (drop (block $label0 (result (ref $s))
+      (return (br_on_cast_fail $label0 (ref $s) (ref $t) (local.get $x)))))
+    (struct.new $t))
+  )
diff --git a/test/lit/passes/flatten_br_on_non_null.wast b/test/lit/passes/flatten_br_on_non_null.wast
new file mode 100644
index 00000000000..b6a83205718
--- /dev/null
+++ b/test/lit/passes/flatten_br_on_non_null.wast
@@ -0,0 +1,64 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_passes_tests_to_lit.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $s (sub (struct)))
+ (type $s (sub (struct)))
+ (type $t (sub $s (struct)))
+
+  ;; CHECK:      (type $1 (func (param (ref null $s)) (result (ref $s))))
+
+  ;; CHECK:      (func $br_on_non_null (type $1) (param $x (ref null $s)) (result (ref $s))
+  ;; CHECK-NEXT:  (local $1 (ref null $s))
+  ;; CHECK-NEXT:  (local $2 (ref $s))
+  ;; CHECK-NEXT:  (local $3 (ref null $s))
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 i32)
+  ;; CHECK-NEXT:  (local $6 (ref (exact $s)))
+  ;; CHECK-NEXT:  (local $7 (ref null $s))
+  ;; CHECK-NEXT:  (local $8 (ref $s))
+  ;; CHECK-NEXT:  (block $label0
+  ;; CHECK-NEXT:   (local.set $1
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $3
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $4
+  ;; CHECK-NEXT:    (ref.is_null
+  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $5
+  ;; CHECK-NEXT:    (i32.eqz
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (br_if $label0
+  ;; CHECK-NEXT:    (local.get $5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:    (struct.new_default $s)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $7
+  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $8
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $7)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return
+  ;; CHECK-NEXT:   (local.get $8)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $br_on_non_null (param $x (ref null $s)) (result (ref $s))
+    (return (block $label0 (result (ref $s))
+      (br_on_non_null $label0 (local.get $x))
+      (struct.new $s))))
+)
+
diff --git a/test/lit/passes/flatten_br_on_null.wast b/test/lit/passes/flatten_br_on_null.wast
new file mode 100644
index 00000000000..d9423bf1c1d
--- /dev/null
+++ b/test/lit/passes/flatten_br_on_null.wast
@@ -0,0 +1,69 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_passes_tests_to_lit.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
+
+(module
+ ;; CHECK:      (type $s (sub (struct)))
+ (type $s (sub (struct)))
+ (type $t (sub $s (struct)))
+
+  ;; CHECK:      (type $1 (func (param (ref null $s)) (result (ref $s))))
+
+  ;; CHECK:      (func $br_on_null (type $1) (param $x (ref null $s)) (result (ref $s))
+  ;; CHECK-NEXT:  (local $1 (ref null $s))
+  ;; CHECK-NEXT:  (local $2 (ref $s))
+  ;; CHECK-NEXT:  (local $3 (ref null $s))
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 (ref $s))
+  ;; CHECK-NEXT:  (local $6 (ref (exact $s)))
+  ;; CHECK-NEXT:  (local $7 (ref $s))
+  ;; CHECK-NEXT:  (local $8 (ref $s))
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (block $label0
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (local.get $1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $4
+  ;; CHECK-NEXT:     (ref.is_null
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (br_if $label0
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $5
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (local.get $5)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:    (struct.new_default $s)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $7
+  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $8
+  ;; CHECK-NEXT:   (local.get $7)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return
+  ;; CHECK-NEXT:   (local.get $8)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $br_on_null (param $x (ref null $s)) (result (ref $s))
+    (block $label0
+      (return (br_on_null $label0 (local.get $x))))
+    (struct.new $s))
+  )

From 836724fff123b0e797603304bec03cca63198fde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 09:46:41 +0100
Subject: [PATCH 07/11] Remove temp file

---
 test.wast | 24 ------------------------
 1 file changed, 24 deletions(-)
 delete mode 100644 test.wast

diff --git a/test.wast b/test.wast
deleted file mode 100644
index 260cfcec323..00000000000
--- a/test.wast
+++ /dev/null
@@ -1,24 +0,0 @@
-(module
- (type $s (sub (struct)))
- (type $t (sub $s (struct)))
-
-  (func $br_on_null (param $x (ref null $s)) (result (ref $s))
-    (block $label0
-      (return (br_on_null $label0 (local.get $x))))
-    (struct.new $s))
-
-  (func $br_on_non_null (param $x (ref null $s)) (result (ref $s))
-    (return (block $label0 (result (ref $s))
-      (br_on_non_null $label0 (local.get $x))
-      (struct.new $s))))
-
-  (func $br_on_cast (param $x (ref $s)) (result (ref $t))
-    (return (block $label0 (result (ref $t))
-      (drop (br_on_cast $label0 (ref $s) (ref $t) (local.get $x)))
-      (struct.new $t))))
-
-  (func $br_on_cast_fail (param $x (ref $s)) (result (ref $t))
-    (drop (block $label0 (result (ref $s))
-      (return (br_on_cast_fail $label0 (ref $s) (ref $t) (local.get $x)))))
-    (struct.new $t))
-  )

From 1bd3cc0ad0924432ff908b7d882164806dcad715 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 09:56:12 +0100
Subject: [PATCH 08/11] Simplify

---
 src/passes/Flatten.cpp | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index e47e8c05cfc..ff0808c24b6 100644
--- a/src/passes/Flatten.cpp
+++ b/src/passes/Flatten.cpp
@@ -311,9 +311,6 @@ struct Flatten
       } else if (auto* br = curr->dynCast<BrOn>()) {
         if (br->op == BrOnOp::BrOnNull || br->op == BrOnOp::BrOnNonNull) {
           auto nullableType = br->ref->type;
-          auto nonNullableType = nullableType.with(Nullability::NonNullable);
-          Index nonNullableTemp =
-            builder.addVar(getFunction(), nonNullableType);
 
           Index nullableTemp = builder.addVar(getFunction(), nullableType);
           ourPreludes.push_back(builder.makeLocalSet(nullableTemp, br->ref));
@@ -328,14 +325,9 @@ struct Flatten
             ourPreludes.push_back(builder.makeBreak(
               br->name, nullptr, builder.makeLocalGet(isNullTemp, Type::i32)));
 
-            ourPreludes.push_back(builder.makeLocalSet(
-              nonNullableTemp,
-              builder.makeRefAs(
-                RefAsOp::RefAsNonNull,
-                builder.makeLocalGet(nullableTemp, nullableType))));
-
-            replaceCurrent(
-              builder.makeLocalGet(nonNullableTemp, nonNullableType));
+            replaceCurrent(builder.makeRefAs(
+              RefAsOp::RefAsNonNull,
+              builder.makeLocalGet(nullableTemp, nullableType)));
           } else { // br_on_non_null
             Index isNotNullTemp = builder.addVar(getFunction(), Type::i32);
             ourPreludes.push_back(builder.makeLocalSet(

From 8f2076dffe80b5f3a81a6f110b06ca7c66801946 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 10:26:16 +0100
Subject: [PATCH 09/11] Update tests

---
 test/lit/passes/flatten_br_on_non_null.wast | 35 +++++++++---------
 test/lit/passes/flatten_br_on_null.wast     | 40 ++++++++++-----------
 2 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/test/lit/passes/flatten_br_on_non_null.wast b/test/lit/passes/flatten_br_on_non_null.wast
index b6a83205718..1df087b27f8 100644
--- a/test/lit/passes/flatten_br_on_non_null.wast
+++ b/test/lit/passes/flatten_br_on_non_null.wast
@@ -12,47 +12,46 @@
 
   ;; CHECK:      (func $br_on_non_null (type $1) (param $x (ref null $s)) (result (ref $s))
   ;; CHECK-NEXT:  (local $1 (ref null $s))
-  ;; CHECK-NEXT:  (local $2 (ref $s))
-  ;; CHECK-NEXT:  (local $3 (ref null $s))
+  ;; CHECK-NEXT:  (local $2 (ref null $s))
+  ;; CHECK-NEXT:  (local $3 i32)
   ;; CHECK-NEXT:  (local $4 i32)
-  ;; CHECK-NEXT:  (local $5 i32)
-  ;; CHECK-NEXT:  (local $6 (ref (exact $s)))
-  ;; CHECK-NEXT:  (local $7 (ref null $s))
-  ;; CHECK-NEXT:  (local $8 (ref $s))
+  ;; CHECK-NEXT:  (local $5 (ref (exact $s)))
+  ;; CHECK-NEXT:  (local $6 (ref null $s))
+  ;; CHECK-NEXT:  (local $7 (ref $s))
   ;; CHECK-NEXT:  (block $label0
   ;; CHECK-NEXT:   (local.set $1
   ;; CHECK-NEXT:    (local.get $x)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $3
+  ;; CHECK-NEXT:   (local.set $2
   ;; CHECK-NEXT:    (local.get $1)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $4
+  ;; CHECK-NEXT:   (local.set $3
   ;; CHECK-NEXT:    (ref.is_null
-  ;; CHECK-NEXT:     (local.get $3)
+  ;; CHECK-NEXT:     (local.get $2)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $5
+  ;; CHECK-NEXT:   (local.set $4
   ;; CHECK-NEXT:    (i32.eqz
-  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:     (local.get $3)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:   (br_if $label0
-  ;; CHECK-NEXT:    (local.get $5)
+  ;; CHECK-NEXT:    (local.get $4)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:   (local.set $5
   ;; CHECK-NEXT:    (struct.new_default $s)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $7
-  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:    (local.get $5)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (local.set $8
+  ;; CHECK-NEXT:  (local.set $7
   ;; CHECK-NEXT:   (ref.as_non_null
-  ;; CHECK-NEXT:    (local.get $7)
+  ;; CHECK-NEXT:    (local.get $6)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (return
-  ;; CHECK-NEXT:   (local.get $8)
+  ;; CHECK-NEXT:   (local.get $7)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (unreachable)
   ;; CHECK-NEXT: )
diff --git a/test/lit/passes/flatten_br_on_null.wast b/test/lit/passes/flatten_br_on_null.wast
index d9423bf1c1d..83f76139a4b 100644
--- a/test/lit/passes/flatten_br_on_null.wast
+++ b/test/lit/passes/flatten_br_on_null.wast
@@ -12,54 +12,50 @@
 
   ;; CHECK:      (func $br_on_null (type $1) (param $x (ref null $s)) (result (ref $s))
   ;; CHECK-NEXT:  (local $1 (ref null $s))
-  ;; CHECK-NEXT:  (local $2 (ref $s))
-  ;; CHECK-NEXT:  (local $3 (ref null $s))
-  ;; CHECK-NEXT:  (local $4 i32)
-  ;; CHECK-NEXT:  (local $5 (ref $s))
-  ;; CHECK-NEXT:  (local $6 (ref (exact $s)))
+  ;; CHECK-NEXT:  (local $2 (ref null $s))
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 (ref $s))
+  ;; CHECK-NEXT:  (local $5 (ref (exact $s)))
+  ;; CHECK-NEXT:  (local $6 (ref $s))
   ;; CHECK-NEXT:  (local $7 (ref $s))
-  ;; CHECK-NEXT:  (local $8 (ref $s))
   ;; CHECK-NEXT:  (block
   ;; CHECK-NEXT:   (block $label0
   ;; CHECK-NEXT:    (local.set $1
   ;; CHECK-NEXT:     (local.get $x)
   ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:    (local.set $2
   ;; CHECK-NEXT:     (local.get $1)
   ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $4
+  ;; CHECK-NEXT:    (local.set $3
   ;; CHECK-NEXT:     (ref.is_null
-  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:      (local.get $2)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (br_if $label0
-  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:     (local.get $3)
   ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:    (local.set $4
   ;; CHECK-NEXT:     (ref.as_non_null
-  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:      (local.get $2)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:    (local.set $5
-  ;; CHECK-NEXT:     (local.get $2)
-  ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (return
-  ;; CHECK-NEXT:     (local.get $5)
+  ;; CHECK-NEXT:     (local.get $4)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (unreachable)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:   (local.set $5
   ;; CHECK-NEXT:    (struct.new_default $s)
   ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (local.set $7
-  ;; CHECK-NEXT:    (local.get $6)
+  ;; CHECK-NEXT:   (local.set $6
+  ;; CHECK-NEXT:    (local.get $5)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (local.set $8
-  ;; CHECK-NEXT:   (local.get $7)
+  ;; CHECK-NEXT:  (local.set $7
+  ;; CHECK-NEXT:   (local.get $6)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (return
-  ;; CHECK-NEXT:   (local.get $8)
+  ;; CHECK-NEXT:   (local.get $7)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $br_on_null (param $x (ref null $s)) (result (ref $s))

From dbd78081fdec7575c20304462843fa964b24b221 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Thu, 2 Oct 2025 11:16:25 +0100
Subject: [PATCH 10/11] Fix br_on_non_null value

---
 src/passes/Flatten.cpp                      | 19 ++++++++++++++----
 test/lit/passes/flatten_br_on_non_null.wast | 22 ++++++++++++++-------
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/passes/Flatten.cpp b/src/passes/Flatten.cpp
index ff0808c24b6..38020e6a3b8 100644
--- a/src/passes/Flatten.cpp
+++ b/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;
diff --git a/test/lit/passes/flatten_br_on_non_null.wast b/test/lit/passes/flatten_br_on_non_null.wast
index 1df087b27f8..831763d357b 100644
--- a/test/lit/passes/flatten_br_on_non_null.wast
+++ b/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

From 4ce98a155e297c29c87214459d7c53b307f21c52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20A=C4=9Facan?= <omersa@google.com>
Date: Fri, 3 Oct 2025 09:40:31 +0100
Subject: [PATCH 11/11] Test formatting

---
 test/lit/passes/flatten_br_on_cast.wast      | 21 ++++++++++++-------
 test/lit/passes/flatten_br_on_cast_fail.wast | 22 +++++++++++++-------
 test/lit/passes/flatten_br_on_non_null.wast  | 18 ++++++++++------
 test/lit/passes/flatten_br_on_null.wast      | 14 ++++++++-----
 4 files changed, 50 insertions(+), 25 deletions(-)

diff --git a/test/lit/passes/flatten_br_on_cast.wast b/test/lit/passes/flatten_br_on_cast.wast
index dc267584550..e8308ef5e6a 100644
--- a/test/lit/passes/flatten_br_on_cast.wast
+++ b/test/lit/passes/flatten_br_on_cast.wast
@@ -4,10 +4,10 @@
 ;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
 
 (module
- ;; CHECK:      (type $s (sub (struct)))
- (type $s (sub (struct)))
- ;; CHECK:      (type $t (sub $s (struct)))
- (type $t (sub $s (struct)))
+  ;; CHECK:      (type $s (sub (struct)))
+  (type $s (sub (struct)))
+  ;; CHECK:      (type $t (sub $s (struct)))
+  (type $t (sub $s (struct)))
 
   ;; CHECK:      (type $2 (func (param (ref $s)) (result (ref $t))))
 
@@ -70,8 +70,15 @@
   ;; CHECK-NEXT:  (unreachable)
   ;; CHECK-NEXT: )
   (func $br_on_cast (param $x (ref $s)) (result (ref $t))
-    (return (block $label0 (result (ref $t))
-      (drop (br_on_cast $label0 (ref $s) (ref $t) (local.get $x)))
-      (struct.new $t))))
+    (return
+      (block $label0 (result (ref $t))
+        (drop
+          (br_on_cast $label0 (ref $s) (ref $t)
+            (local.get $x)
+          )
+        )
+      (struct.new $t))
+    )
+  )
 )
 
diff --git a/test/lit/passes/flatten_br_on_cast_fail.wast b/test/lit/passes/flatten_br_on_cast_fail.wast
index 2c55763a615..8c13eed805a 100644
--- a/test/lit/passes/flatten_br_on_cast_fail.wast
+++ b/test/lit/passes/flatten_br_on_cast_fail.wast
@@ -4,10 +4,10 @@
 ;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
 
 (module
- ;; CHECK:      (type $s (sub (struct)))
- (type $s (sub (struct)))
- ;; CHECK:      (type $t (sub $s (struct)))
- (type $t (sub $s (struct)))
+  ;; CHECK:      (type $s (sub (struct)))
+  (type $s (sub (struct)))
+  ;; CHECK:      (type $t (sub $s (struct)))
+  (type $t (sub $s (struct)))
 
   ;; CHECK:      (type $2 (func (param (ref $s)) (result (ref $t))))
 
@@ -86,7 +86,15 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $br_on_cast_fail (param $x (ref $s)) (result (ref $t))
-    (drop (block $label0 (result (ref $s))
-      (return (br_on_cast_fail $label0 (ref $s) (ref $t) (local.get $x)))))
-    (struct.new $t))
+    (drop
+      (block $label0 (result (ref $s))
+        (return
+          (br_on_cast_fail $label0 (ref $s) (ref $t)
+            (local.get $x)
+          )
+        )
+      )
+    )
+    (struct.new $t)
   )
+)
diff --git a/test/lit/passes/flatten_br_on_non_null.wast b/test/lit/passes/flatten_br_on_non_null.wast
index 831763d357b..b58d6a49f23 100644
--- a/test/lit/passes/flatten_br_on_non_null.wast
+++ b/test/lit/passes/flatten_br_on_non_null.wast
@@ -4,9 +4,9 @@
 ;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
 
 (module
- ;; CHECK:      (type $s (sub (struct)))
- (type $s (sub (struct)))
- (type $t (sub $s (struct)))
+  ;; CHECK:      (type $s (sub (struct)))
+  (type $s (sub (struct)))
+  (type $t (sub $s (struct)))
 
   ;; CHECK:      (type $1 (func (param (ref null $s)) (result (ref $s))))
 
@@ -64,8 +64,14 @@
   ;; CHECK-NEXT:  (unreachable)
   ;; CHECK-NEXT: )
   (func $br_on_non_null (param $x (ref null $s)) (result (ref $s))
-    (return (block $label0 (result (ref $s))
-      (br_on_non_null $label0 (local.get $x))
-      (struct.new $s))))
+    (return
+      (block $label0 (result (ref $s))
+        (br_on_non_null $label0
+          (local.get $x)
+        )
+        (struct.new $s)
+      )
+    )
+  )
 )
 
diff --git a/test/lit/passes/flatten_br_on_null.wast b/test/lit/passes/flatten_br_on_null.wast
index 83f76139a4b..1edc253f11f 100644
--- a/test/lit/passes/flatten_br_on_null.wast
+++ b/test/lit/passes/flatten_br_on_null.wast
@@ -4,9 +4,9 @@
 ;; RUN: foreach %s %t wasm-opt --flatten --all-features -S -o - | filecheck %s
 
 (module
- ;; CHECK:      (type $s (sub (struct)))
- (type $s (sub (struct)))
- (type $t (sub $s (struct)))
+  ;; CHECK:      (type $s (sub (struct)))
+  (type $s (sub (struct)))
+  (type $t (sub $s (struct)))
 
   ;; CHECK:      (type $1 (func (param (ref null $s)) (result (ref $s))))
 
@@ -60,6 +60,10 @@
   ;; CHECK-NEXT: )
   (func $br_on_null (param $x (ref null $s)) (result (ref $s))
     (block $label0
-      (return (br_on_null $label0 (local.get $x))))
-    (struct.new $s))
+      (return
+        (br_on_null $label0 (local.get $x))
+      )
+    )
+    (struct.new $s)
   )
+)