/*
* Copyright 2023 WebAssembly Community Group participants
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
//
// Refine types based on global information about abstract types, that is, types
// that are not created anywhere (no struct.new etc.).
//
// In trapsNeverHappen mode, if we see a cast to $B and the type hierarchy is
// this:
//
// $A :> $B :> $C
//
// and $B has no struct.new instructions, and we are in closed world, then we
// can infer that the cast must be to $C. That is necessarily so since we will
// not trap by assumption, and $C or a subtype of it is all that remains
// possible.
//
// Even without trapsNeverHappen we can optimize certain cases. When we see a
// cast to a type that is never created, nor any subtype is created, then it
// must fail unless it allows null.
//
#include <memory>
#include "ir/drop.h"
#include "ir/localize.h"
#include "ir/module-utils.h"
#include "ir/subtypes.h"
#include "ir/type-updating.h"
#include "ir/utils.h"
#include "pass.h"
#include "wasm-type.h"
#include "wasm.h"
namespace wasm {
namespace {
using Types = std::unordered_set<HeapType>;
// Gather all types in StructNews.
struct NewFinder : public PostWalker<NewFinder> {
Types& types;
NewFinder(Types& types) : types(types) {}
void visitStructNew(StructNew* curr) {
auto type = curr->type;
if (type != Type::unreachable) {
types.insert(type.getHeapType());
}
}
};
struct AbstractTypeRefining : public Pass {
// Changes types by refining them. We never add new non-nullable locals here
// (even if we refine a type to a bottom type, we only change the heap type
// there, not nullability).
bool requiresNonNullableLocalFixups() override { return false; }
// The types that are created (have a struct.new).
Types createdTypes;
// The types that are created, or have a subtype that is created.
Types createdTypesOrSubTypes;
// A map of a cast type to refine and the type to refine it to.
TypeMapper::TypeUpdates refinableTypes;
bool trapsNeverHappen;
void run(Module* module) override {
if (!module->features.hasGC()) {
return;
}
if (getPassOptions().worldMode == WorldMode::Open) {
Fatal() << "AbstractTypeRefining requires --closed-world";
}
trapsNeverHappen = getPassOptions().trapsNeverHappen;
// First, find all the created types (that have a struct.new) both in module
// code and in functions.
NewFinder(createdTypes).walkModuleCode(module);
ModuleUtils::ParallelFunctionAnalysis<Types> analysis(
*module, [&](Function* func, Types& types) {
if (!func->imported()) {
NewFinder(types).walk(func->body);
}
});
for (auto& [_, types] : analysis.map) {
for (auto type : types) {
createdTypes.insert(type);
}
}
// Assume all public types are created, which makes them non-abstract and
// hence ignored below.
// TODO: In principle we could assume such types are not created outside the
// module, given closed world, but we'd also need to make sure that
// we don't need to make any changes to public types that refer to
// them.
for (auto type :
ModuleUtils::getPublicHeapTypes(*module, getPassOptions().worldMode)) {
createdTypes.insert(type);
}
SubTypes subTypes(*module);
// Compute createdTypesOrSubTypes by starting with the created types and
// then propagating subtypes.
createdTypesOrSubTypes = createdTypes;
for (auto type : subTypes.getSubTypesFirstSort()) {
// If any of our subtypes are created, so are we.
for (auto subType : subTypes.getImmediateSubTypes(type)) {
if (createdTypesOrSubTypes.contains(subType)) {
createdTypesOrSubTypes.insert(type);
break;
}
}
}
if (trapsNeverHappen) {
computeAbstractTypes(subTypes);
}
// Use what we found about abstract types and never-created types to
// optimize.
optimize(module, subTypes);
}
void computeAbstractTypes(const SubTypes& subTypes) {
// Abstract types are those with no news, i.e., the complement of
// |createdTypes|. As mentioned above, we can only optimize this case if
// traps never happen.
// TODO: We could do some of this even if traps are possible. If an abstract
// type has no casts at all, then no traps are relevant, and we could
// remove it from the module. That might also make sense in MergeTypes
// perhaps (which atm will not merge such types if they add fields,
// in particular).
Types abstractTypes;
for (auto type : subTypes.types) {
if (!createdTypes.contains(type)) {
abstractTypes.insert(type);
}
}
// We found abstract types. Next, find which of them are refinable. We
// need an abstract type to have a single subtype, to which we will switch
// all of their casts.
//
// Do this depth-first, so that we visit subtypes first. That will handle
// chains where we want to refine a type A to a subtype of a subtype of
// it.
for (auto type : subTypes.getSubTypesFirstSort()) {
if (!abstractTypes.contains(type)) {
continue;
}
std::optional<HeapType> refinedType;
auto& typeSubTypes = subTypes.getImmediateSubTypes(type);
if (typeSubTypes.size() == 1) {
// There is only a single possibility, so we can definitely use that
/// one.
refinedType = typeSubTypes[0];
} else if (!typeSubTypes.empty()) {
// There are multiple possibilities. However, perhaps only one of them
// is relevant, if nothing is ever created of the others or their
// subtypes.
for (auto subType : typeSubTypes) {
if (createdTypesOrSubTypes.contains(subType)) {
if (!refinedType) {
// This is the first relevant thing, and hopefully will remain
// the only one.
refinedType = subType;
} else {
// We've seen more than one as relevant, so we have failed to
// find a singleton.
refinedType = std::nullopt;
break;
}
}
}
}
if (refinedType) {
// Propagate anything from the child, to handle chains.
auto iter = refinableTypes.find(*refinedType);
if (iter != refinableTypes.end()) {
*refinedType = iter->second;
}
refinableTypes[type] = *refinedType;
}
}
}
void optimize(Module* module, const SubTypes& subTypes) {
// To optimize we rewrite types. That is, if we want to optimize all casts
// of $A to instead cast to the refined type $B, we can do that by simply
// replacing all appearances of $A with $B. That is possible here since we
// only optimize when we know $A is never created, and we are removing all
// casts to it, which means no other references to it are needed - so we can
// just rewrite all references to $A to point to $B. Doing such a rewrite
// will also remove the unneeded type from the type section, which is nice
// for code size.
//
// Even though this pass removes types, it does not on its own inhibit
// further optimizations. In more detail, a possible issue could have been
// something like this: imagine that we replace all $A with $B, and we had
// types like this:
//
// $C = [.., $A, ..]
// $D = [.., $B, ..]
//
// After replacing $A with $B, we cause $C and $D to be structurally
// identical. If we merged $C and $D then we might lose some optimization
// potential (perhaps different values are written to each, and GUFA or
// another pass can optimize each separately, but not if they were merged).
// However, the type rewriter will create a single new rec group for all new
// types anyhow, so they all remain distinct from each other. The only thing
// that would actually merge them is if we run TypeMerging, which is not run
// by default exactly for this reason, that it can limit optimizations.
// Thus, this pass does only "safe" merging, that cannot limit later
// optimizations - merging $A and $B is of course fine as one of them was
// not even used anywhere.
TypeMapper::TypeUpdates mapping;
for (auto type : subTypes.types) {
if (!type.isStruct()) {
// TODO: Support arrays and functions (for functions we will need to
// handle configureAll).
continue;
}
// Add a mapping of types that are never created (and none of their
// subtypes) to the bottom type. This is valid because all locations of
// that type, like a local variable, will only contain null at runtime.
// Likewise, if we have a ref.test of such a type, we can only be looking
// for a null at best. This can be seen as "refining" uses of these
// never-created types to the bottom type.
//
// We check this first as it is the most powerful change.
if (!createdTypesOrSubTypes.contains(type)) {
mapping[type] = type.getBottom();
continue;
}
// Otherwise, apply a refining if we found one before.
if (auto iter = refinableTypes.find(type); iter != refinableTypes.end()) {
mapping[type] = iter->second;
}
}
if (mapping.empty()) {
return;
}
// We may need to apply some preliminary optimizations to ensure we maintain
// validity and correctness when we rewrite the types.
preoptimize(*module, mapping);
// Rewriting types can usually rewrite subtype relationships. For example,
// if we have this:
//
// C <: B <: A
//
// And we see that B is never created, we would naively map B to its subtype
// C. But if we rewrote C's supertype, C would declare itself to be its own
// supertype, which is not allowed. We could fix this by walking up the
// supertype chain to find a supertype that is not being rewritten, but
// changing subtype relationships and keeping descriptor chains valid is
// nontrivial. Instead, avoid changing subtype relationships entirely: leave
// that for Unsubtyping.
class AbstractTypeRefiningTypeMapper : public TypeMapper {
public:
AbstractTypeRefiningTypeMapper(Module& wasm,
const TypeUpdates& mapping,
WorldMode worldMode)
: TypeMapper(wasm, mapping, worldMode) {}
std::optional<HeapType> getDeclaredSuperType(HeapType oldType) override {
// We do not want to update subtype relationships.
return oldType.getDeclaredSuperType();
}
};
AbstractTypeRefiningTypeMapper(*module, mapping, getPassOptions().worldMode)
.map();
// Refinalize to propagate the type changes we made. For example, a refined
// cast may lead to a struct.get reading a more refined type using that
// type.
ReFinalize().run(getPassRunner(), module);
}
void preoptimize(Module& module, const TypeMapper::TypeUpdates& mapping) {
if (!module.features.hasCustomDescriptors()) {
// No descriptor casts, exact casts, or allocations with descriptors to
// fix up.
return;
}
struct Preoptimizer : WalkerPass<PostWalker<Preoptimizer>> {
const TypeMapper::TypeUpdates& mapping;
Preoptimizer(const TypeMapper::TypeUpdates& mapping) : mapping(mapping) {}
bool isFunctionParallel() override { return true; }
std::unique_ptr<Pass> create() override {
return std::make_unique<Preoptimizer>(mapping);
}
Block* localizeChildren(Expression* curr) {
return ChildLocalizer(
curr, getFunction(), *getModule(), getPassOptions())
.getChildrenReplacement();
}
std::optional<HeapType> getOptimized(Type type) {
if (!type.isRef()) {
return std::nullopt;
}
auto heapType = type.getHeapType();
auto it = mapping.find(heapType);
if (it == mapping.end()) {
return std::nullopt;
}
assert(it->second != heapType);
return it->second;
}
// We may have casts like this:
//
// (ref.cast_desc_eq (ref null $optimized-to-bottom)
// (some struct...)
// (some desc...)
// )
//
// We will optimize the cast target to nullref, but then ReFinalize would
// fix up the cast target back to $optimized-to-bottom. Optimize out the
// cast (which we know must either be a null check or unconditional trap)
// to avoid this reintroduction of the optimized type.
//
// Separately, we may have exact casts like this:
//
// (br_on_cast anyref $l (ref (exact $uninstantiated)) ... )
//
// We know such casts will fail (or will pass only for null values), but
// with traps-never-happen, we might optimize them to this:
//
// (br_on_cast anyref $l (ref (exact $instantiated-subtype)) ... )
//
// This might cause the casts to incorrectly start succeeding. To avoid
// that, optimize them out first.
void visitRefCast(RefCast* curr) {
auto optimized = getOptimized(curr->type);
if (!optimized) {
return;
}
// Exact casts to any optimized type and descriptor casts whose types
// will be optimized to bottom either admit null or fail
// unconditionally. Optimize to a cast to bottom, reusing curr and
// preserving nullability. We may need to move the ref value past the
// descriptor value, if any.
Builder builder(*getModule());
if (curr->type.isExact() || (curr->desc && optimized->isBottom())) {
if (curr->desc) {
if (curr->desc->type.isNullableRef() &&
!getPassOptions().trapsNeverHappen) {
curr->desc = builder.makeRefAs(RefAsNonNull, curr->desc);
}
Block* replacement = localizeChildren(curr);
curr->desc = nullptr;
curr->type = curr->type.with(optimized->getBottom());
replacement->list.push_back(curr);
replacement->type = curr->type;
replaceCurrent(replacement);
} else {
curr->type = curr->type.with(optimized->getBottom());
}
}
}
// If we optimize away a descriptor type, then we must fix up any
// ref.get_desc of it, as ReFinalize would fix us up to return it.
void visitRefGetDesc(RefGetDesc* curr) {
auto optimized = getOptimized(curr->type);
if (!optimized) {
return;
}
// We are optimizing the descriptor to a subtype.
auto subDescriptor = *optimized;
Builder builder(*getModule());
if (curr->type.isExact() ||
subDescriptor.isMaybeShared(HeapType::none)) {
// This is exact, so we can ignore subtypes, or there is no subtype to
// optimize to. In this case it must trap.
replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
builder.makeUnreachable()));
} else {
// The descriptor is abstract, but it has a subtype that is not
// (which we want to optimize to). We know this will only succeed if
// we receive the subtype of the described type, so that that
// descriptor is fetched. Add a cast so that we validate.
auto subDescribed = subDescriptor.getDescribedType();
assert(subDescribed);
curr->ref =
builder.makeRefCast(curr->ref, curr->ref->type.with(*subDescribed));
}
}
void visitBrOn(BrOn* curr) {
if (curr->op == BrOnNull || curr->op == BrOnNonNull) {
return;
}
auto optimized = getOptimized(curr->castType);
if (!optimized) {
return;
}
// Optimize the same way we optimize ref.cast*.
Builder builder(*getModule());
bool isFail = curr->op == BrOnCastDescEqFail;
if (curr->castType.isExact() || (curr->desc && optimized->isBottom())) {
if (curr->desc) {
if (curr->desc->type.isNullableRef() &&
!getPassOptions().trapsNeverHappen) {
curr->desc = builder.makeRefAs(RefAsNonNull, curr->desc);
}
Block* replacement = localizeChildren(curr);
// Reuse `curr` as a br_on_cast to nullref. Leave further
// optimization of the branch to RemoveUnusedBrs.
curr->desc = nullptr;
curr->castType = curr->castType.with(optimized->getBottom());
if (isFail) {
curr->op = BrOnCastFail;
curr->type = curr->castType;
} else {
curr->op = BrOnCast;
}
replacement->list.push_back(curr);
replacement->type = curr->type;
replaceCurrent(replacement);
} else {
curr->castType = curr->castType.with(optimized->getBottom());
}
}
}
void visitStructNew(StructNew* curr) {
if (!curr->desc) {
return;
}
auto optimized = getOptimized(curr->desc->type);
if (!optimized) {
return;
}
// The descriptor type is not instantiated, so there is no way this
// allocation can succeed. We need to remove it to avoid leaving it with
// a mismatched descriptor type after type rewriting. If we are in a
// function context, replace it with unreachable, taking care to
// preserve any side effects. If we're not in a function context, then
// we cannot use things like blocks or drops, but there are also no
// effects besides traps, so we can just replace the descriptor with a
// null.
Builder builder(*getModule());
if (getFunction()) {
replaceCurrent(getDroppedChildrenAndAppend(
curr, *getModule(), getPassOptions(), builder.makeUnreachable()));
} else {
curr->desc = builder.makeRefNull(HeapType::none);
}
}
} preoptimizer(mapping);
preoptimizer.run(getPassRunner(), &module);
preoptimizer.runOnModuleCode(getPassRunner(), &module);
}
};
} // anonymous namespace
Pass* createAbstractTypeRefiningPass() { return new AbstractTypeRefining(); }
} // namespace wasm