blob: 2b45b73c9281a926328d6c31a4e0332801a0d7fb [file] [log] [blame] [edit]
// Copyright 2022 The Tint Authors.
//
// 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.
#include "src/tint/transform/merge_return.h"
#include <utility>
#include "src/tint/program_builder.h"
#include "src/tint/sem/statement.h"
#include "src/tint/utils/scoped_assignment.h"
TINT_INSTANTIATE_TYPEINFO(tint::transform::MergeReturn);
using namespace tint::number_suffixes; // NOLINT
namespace tint::transform {
namespace {
/// Returns `true` if `stmt` has the behavior `behavior`.
bool HasBehavior(const Program* program, const ast::Statement* stmt, sem::Behavior behavior) {
return program->Sem().Get(stmt)->Behaviors().Contains(behavior);
}
/// Returns `true` if `func` needs to be transformed.
bool NeedsTransform(const Program* program, const ast::Function* func) {
// Entry points and intrinsic declarations never need transforming.
if (func->IsEntryPoint() || func->body == nullptr) {
return false;
}
// Avoid transforming functions that only have a single exit point.
// TODO(jrprice): Alternatively, use the uniformity analysis to decide which
// functions need to be transformed.
for (auto* s : func->body->statements) {
// Find the first statement that has contains the Return behavior.
if (HasBehavior(program, s, sem::Behavior::kReturn)) {
// If this statement is itself a return, it will be the only exit point,
// so no need to apply the transform to the function.
if (s->Is<ast::ReturnStatement>()) {
return false;
} else {
// Apply the transform in all other cases.
return true;
}
}
}
return false;
}
} // namespace
MergeReturn::MergeReturn() = default;
MergeReturn::~MergeReturn() = default;
namespace {
/// Internal class used to during the transform.
class State {
private:
/// The clone context.
CloneContext& ctx;
/// The program builder.
ProgramBuilder& b;
/// The function.
const ast::Function* function;
/// The symbol for the return flag variable.
Symbol flag;
/// The symbol for the return value variable.
Symbol retval;
/// Tracks whether we are currently inside a loop or switch statement.
bool is_in_loop_or_switch = false;
public:
/// Constructor
/// @param context the clone context
State(CloneContext& context, const ast::Function* func)
: ctx(context), b(*ctx.dst), function(func) {}
/// Process a statement (recursively).
void ProcessStatement(const ast::Statement* stmt) {
if (stmt == nullptr || !HasBehavior(ctx.src, stmt, sem::Behavior::kReturn)) {
return;
}
Switch(
stmt, [&](const ast::BlockStatement* block) { ProcessBlock(block); },
[&](const ast::CaseStatement* c) { ProcessStatement(c->body); },
[&](const ast::ForLoopStatement* f) {
TINT_SCOPED_ASSIGNMENT(is_in_loop_or_switch, true);
ProcessStatement(f->body);
},
[&](const ast::IfStatement* i) {
ProcessStatement(i->body);
ProcessStatement(i->else_statement);
},
[&](const ast::LoopStatement* l) {
TINT_SCOPED_ASSIGNMENT(is_in_loop_or_switch, true);
ProcessStatement(l->body);
},
[&](const ast::ReturnStatement* r) {
utils::Vector<const ast::Statement*, 3> stmts;
// Set the return flag to signal that we have hit a return.
stmts.Push(b.Assign(b.Expr(flag), true));
if (r->value) {
// Set the return value if necessary.
stmts.Push(b.Assign(b.Expr(retval), ctx.Clone(r->value)));
}
if (is_in_loop_or_switch) {
// If we are in a loop or switch statement, break out of it.
stmts.Push(b.Break());
}
ctx.Replace(r, b.Block(std::move(stmts)));
},
[&](const ast::SwitchStatement* s) {
TINT_SCOPED_ASSIGNMENT(is_in_loop_or_switch, true);
for (auto* c : s->body) {
ProcessStatement(c);
}
},
[&](const ast::WhileStatement* w) {
TINT_SCOPED_ASSIGNMENT(is_in_loop_or_switch, true);
ProcessStatement(w->body);
},
[&](Default) { TINT_ICE(Transform, b.Diagnostics()) << "unhandled statement type"; });
}
void ProcessBlock(const ast::BlockStatement* block) {
// We will rebuild the contents of the block statement.
// We may introduce conditionals around statements that follow a statement with the
// `Return` behavior, so build a stack of statement lists that represent the new
// (potentially nested) conditional blocks.
utils::Vector<utils::Vector<const ast::Statement*, 8>, 8> new_stmts({{}});
// Insert variables for the return flag and return value at the top of the function.
if (block == function->body) {
flag = b.Symbols().New("tint_return_flag");
new_stmts[0].Push(b.Decl(b.Var(flag, b.ty.bool_())));
if (!function->return_type->Is<ast::Void>()) {
retval = b.Symbols().New("tint_return_value");
new_stmts[0].Push(b.Decl(b.Var(retval, ctx.Clone(function->return_type))));
}
}
for (auto* s : block->statements) {
// Process the statement and add it to the current block.
ProcessStatement(s);
new_stmts.Back().Push(ctx.Clone(s));
// Check if the statement is or contains a return statement.
// We need to make sure any statements that follow this one do not get executed if the
// return flag has been set.
if (HasBehavior(ctx.src, s, sem::Behavior::kReturn)) {
if (is_in_loop_or_switch) {
// We're in a loop/switch, and so we would have inserted a `break`.
// If we've just come out of a loop/switch statement, we need to `break` again.
if (s->IsAnyOf<ast::LoopStatement, ast::ForLoopStatement,
ast::SwitchStatement>()) {
// If the loop only has the 'Return' behavior, we can just unconditionally
// break. Otherwise check the return flag.
if (HasBehavior(ctx.src, s, sem::Behavior::kNext)) {
new_stmts.Back().Push(
b.If(b.Expr(flag), b.Block(utils::Vector{b.Break()})));
} else {
new_stmts.Back().Push(b.Break());
}
}
} else {
// Create a new list for any subsequent statements, which we will wrap in a
// conditional block.
new_stmts.Push({});
}
}
}
// Descend the stack of new block statements, wrapping them in conditionals.
while (new_stmts.Length() > 1) {
const ast::IfStatement* i = nullptr;
if (new_stmts.Back().Length() > 0) {
i = b.If(b.Not(b.Expr(flag)), b.Block(new_stmts.Back()));
}
new_stmts.Pop();
if (i) {
new_stmts.Back().Push(i);
}
}
// Insert the final return statement at the end of the function body.
if (block == function->body && retval.IsValid()) {
new_stmts[0].Push(b.Return(b.Expr(retval)));
}
ctx.Replace(block, b.Block(new_stmts[0]));
}
};
} // namespace
Transform::ApplyResult MergeReturn::Apply(const Program* src, const DataMap&, DataMap&) const {
ProgramBuilder b;
CloneContext ctx{&b, src, /* auto_clone_symbols */ true};
bool made_changes = false;
for (auto* func : ctx.src->AST().Functions()) {
if (!NeedsTransform(ctx.src, func)) {
continue;
}
State state(ctx, func);
state.ProcessStatement(func->body);
made_changes = true;
}
if (!made_changes) {
return SkipTransform;
}
ctx.Clone();
return Program(std::move(b));
}
} // namespace tint::transform