Skip to content

Commit

Permalink
Heap2Local: Properly handle failing array casts (WebAssembly#6772)
Browse files Browse the repository at this point in the history 
Followup to WebAssembly#6727 which added support for failing casts in Struct2Local, but it
turns out that it required Array2Struct changes as well. Specifically, when we
turn an array into a struct then casts can look like they behave differently
(what used to be an array input, becomes a struct), so like with RefTest that we
already handled, check if the cast succeeds in the original form and handle
that.
  • Loading branch information
@kripken
kripken authored Jul 18, 2024
1 parent 84daeca commit a8066e6
Show file tree
Hide file tree
Showing 2 changed files with 164 additions and 5 deletions.
44 changes: 39 additions & 5 deletions src/passes/Heap2Local.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -862,6 +862,11 @@ struct Array2Struct : PostWalker<Array2Struct> {
// The original type of the allocation, before we turn it into a struct.
Type originalType;

// The type of the struct we are changing to (nullable and non-nullable
// variations).
Type nullStruct;
Type nonNullStruct;

Array2Struct(Expression* allocation,
EscapeAnalyzer& analyzer,
Function* func,
Expand Down Expand Up @@ -928,9 +933,15 @@ struct Array2Struct : PostWalker<Array2Struct> {
// lowered away to locals anyhow.
auto nullArray = Type(arrayType, Nullable);
auto nonNullArray = Type(arrayType, NonNullable);
auto nullStruct = Type(structType, Nullable);
auto nonNullStruct = Type(structType, NonNullable);
nullStruct = Type(structType, Nullable);
nonNullStruct = Type(structType, NonNullable);
for (auto* reached : analyzer.reached) {
if (reached->is<RefCast>()) {
// Casts must be handled later: We need to see the old type, and to
// potentially replace the cast based on that, see below.
continue;
}

// We must check subtyping here because the allocation may be upcast as it
// flows around. If we do see such upcasting then we are refining here and
// must refinalize.
Expand Down Expand Up @@ -1032,15 +1043,14 @@ struct Array2Struct : PostWalker<Array2Struct> {
}

// Some additional operations need special handling

void visitRefTest(RefTest* curr) {
if (!analyzer.reached.count(curr)) {
return;
}

// When we ref.test an array allocation, we cannot simply turn the array
// into a struct, as then the test will behave different. (Note that this is
// not a problem for ref.*cast*, as the cast simply goes away when the value
// flows through, and we verify it will do so in the escape analysis.) To
// into a struct, as then the test will behave differently. To properly
// handle this, check if the test succeeds or not, and write out the outcome
// here (similar to Struct2Local::visitRefTest). Note that we test on
// |originalType| here and not |allocation->type|, as the allocation has
Expand All @@ -1050,6 +1060,30 @@ struct Array2Struct : PostWalker<Array2Struct> {
builder.makeConst(Literal(result))));
}

void visitRefCast(RefCast* curr) {
if (!analyzer.reached.count(curr)) {
return;
}

// As with RefTest, we need to check if the cast succeeds with the array
// type before we turn it into a struct type (as after that change, the
// outcome of the cast will look different).
if (!Type::isSubType(originalType, curr->type)) {
// The cast fails, ensure we trap with an unreachable.
replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
builder.makeUnreachable()));
} else {
// The cast succeeds. Update the type. (It is ok to use the non-nullable
// type here unconditionally, since we know the allocation flows through
// here, and anyhow we will be removing the reference during Struct2Local,
// later.)
curr->type = nonNullStruct;
}

// Regardless of how we altered the type here, refinalize.
refinalize = true;
}

// Get the value in an expression we know must contain a constant index.
Index getIndex(Expression* curr) {
return curr->cast<Const>()->value.getUnsigned();
Expand Down
125 changes: 125 additions & 0 deletions test/lit/passes/heap2local.wast
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4303,3 +4303,128 @@
)
)
)

;; Array casts to structs and arrays.
(module
;; CHECK: (type $1 (func))

;; CHECK: (type $0 (func (result (ref struct))))
(type $0 (func (result (ref struct))))
;; CHECK: (type $3 (func (result structref)))

;; CHECK: (type $4 (func (result (ref array))))

;; CHECK: (type $array (array i8))
(type $array (array i8))

;; CHECK: (func $array.cast.struct (type $0) (result (ref struct))
;; CHECK-NEXT: (local $eq (ref eq))
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
(func $array.cast.struct (result (ref struct))
(local $eq (ref eq))
;; This cast will fail: we cast an array to struct. That we go through
;; (ref eq) in the middle, which seems like it could cast to struct, should
;; not confuse us. And, as the cast fails, the reference does not escape.
;; We can optimize here and will emit an unreachable for the failing cast.
(ref.cast (ref struct)
(local.tee $eq
(array.new_fixed $array 0)
)
)
)

;; CHECK: (func $array.cast.struct.null (type $3) (result structref)
;; CHECK-NEXT: (local $eq (ref eq))
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
(func $array.cast.struct.null (result (ref null struct))
(local $eq (ref eq))
;; As above but the cast is to a nullable type, which changes nothing.
(ref.cast (ref null struct)
(local.tee $eq
(array.new_fixed $array 0)
)
)
)

;; CHECK: (func $array.cast.array (type $4) (result (ref array))
;; CHECK-NEXT: (local $eq (ref eq))
;; CHECK-NEXT: (ref.cast (ref array)
;; CHECK-NEXT: (local.tee $eq
;; CHECK-NEXT: (array.new_fixed $array 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $array.cast.array (result (ref array))
(local $eq (ref eq))
;; Now we cast to array, and the cast succeeds, so we escape, and do
;; nothing to optimize.
(ref.cast (ref array)
(local.tee $eq
(array.new_fixed $array 0)
)
)
)

;; CHECK: (func $array.cast.array.set (type $1)
;; CHECK-NEXT: (local $eq (ref eq))
;; CHECK-NEXT: (local $array (ref array))
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $array.cast.array.set
(local $eq (ref eq))
(local $array (ref array))
;; As above, but now we store the result in a local rather than return it
;; out, so it does not escape.
(local.set $array
(ref.cast (ref array)
(local.tee $eq
(array.new_fixed $array 0)
)
)
)
)

;; CHECK: (func $array.cast.struct.set (type $1)
;; CHECK-NEXT: (local $eq (ref eq))
;; CHECK-NEXT: (local $struct (ref struct))
;; CHECK-NEXT: (local.tee $struct
;; CHECK-NEXT: (block
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (unreachable)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $array.cast.struct.set
(local $eq (ref eq))
(local $struct (ref struct))
;; As above, but now the cast fails and is stored to a struct local. We do not
;; escape and we emit an unreachable for the cast.
(local.set $struct
(ref.cast (ref struct)
(local.tee $eq
(array.new_fixed $array 0)
)
)
)
)
)

0 comments on commit a8066e6

Please sign in to comment.