Skip to content

Flatten: Support br_on_null, br_on_non_null, br_on_cast, br_on_cast_fail #7939

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Flatten: Support br_on_null, br_on_non_null, br_on_cast, br_on_cast_fail #7939

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
f8ec112
Implement flattening for br_on_null
osa1 Oct 1, 2025
5983579
Implement flattening for br_on_non_null
osa1 Oct 1, 2025
a0a5bab
Implement br_on_cast
osa1 Oct 1, 2025
e12bf86
WIP
osa1 Oct 1, 2025
1e2420d
Fix bug
osa1 Oct 2, 2025
cd31b32
Add tests
osa1 Oct 2, 2025
836724f
Remove temp file
osa1 Oct 2, 2025
1bd3cc0
Simplify
osa1 Oct 2, 2025
8f2076d
Update tests
osa1 Oct 2, 2025
dbd7808
Fix br_on_non_null value
osa1 Oct 2, 2025
4ce98a1
Test formatting
osa1 Oct 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 114 additions & 1 deletion src/passes/Flatten.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -308,6 +309,118 @@ struct Flatten
}
}

} else if (auto* br = curr->dynCast<BrOn>()) {
if (br->op == BrOnOp::BrOnNull || br->op == BrOnOp::BrOnNonNull) {
auto nullableType = br->ref->type;

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))));

if (br->op == BrOnOp::BrOnNull) {
ourPreludes.push_back(builder.makeBreak(
br->name, nullptr, builder.makeLocalGet(isNullTemp, Type::i32)));

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(
isNotNullTemp,
builder.makeUnary(UnaryOp::EqZInt32,
builder.makeLocalGet(isNullTemp, 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;
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)));

// 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);
}

if (br->op == BrOnCast) {
Index breakTargetTemp = getTempForBreakTarget(br->name, targetType);

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
Index breakTargetTemp =
getTempForBreakTarget(br->name, failValue->type);

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))));

ourPreludes.push_back(
builder.makeIf(builder.makeLocalGet(typeTestFailTemp, Type::i32),
builder.makeBlock(failureBlock)));

Index targetTypeTemp = builder.addVar(getFunction(), targetType);
ourPreludes.push_back(builder.makeLocalSet(
targetTypeTemp,
builder.makeRefCast(
builder.makeLocalGet(sourceTypeTemp, sourceType), targetType)));

replaceCurrent(builder.makeLocalGet(targetTypeTemp, targetType));
}

} else {
Fatal() << "Unsupported instruction for Flatten: BrOn " << br->op;
}
} else if (auto* sw = curr->dynCast<Switch>()) {
if (sw->value) {
auto type = sw->value->type;
Expand All @@ -333,7 +446,7 @@ struct Flatten
}
}

if (curr->is<BrOn>() || curr->is<TryTable>()) {
if (curr->is<TryTable>()) {
Fatal() << "Unsupported instruction for Flatten: "
<< getExpressionName(curr);
}
Expand Down
77 changes: 77 additions & 0 deletions test/lit/passes/flatten_br_on_cast.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))))
)

92 changes: 92 additions & 0 deletions test/lit/passes/flatten_br_on_cast_fail.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
)
71 changes: 71 additions & 0 deletions test/lit/passes/flatten_br_on_non_null.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
;; 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 null $s))
;; CHECK-NEXT: (local $3 i32)
;; CHECK-NEXT: (local $4 i32)
;; 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
;; 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.is_null
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $4
;; CHECK-NEXT: (i32.eqz
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; 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 $6
;; CHECK-NEXT: (struct.new_default $s)
;; CHECK-NEXT: )
;; 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 $5)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (return
;; CHECK-NEXT: (local.get $7)
;; 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))))
)

Loading
Loading