blob: bfe8865ffa5d130531a3dc8712ef180caf5e47eb [file] [log] [blame]
// Copyright 2021 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/localize_struct_array_assignment.h"
#include <unordered_map>
#include <utility>
#include "src/tint/ast/assignment_statement.h"
#include "src/tint/ast/traverse_expressions.h"
#include "src/tint/program_builder.h"
#include "src/tint/sem/expression.h"
#include "src/tint/sem/member_accessor_expression.h"
#include "src/tint/sem/reference.h"
#include "src/tint/sem/statement.h"
#include "src/tint/sem/variable.h"
#include "src/tint/transform/simplify_pointers.h"
#include "src/tint/utils/scoped_assignment.h"
TINT_INSTANTIATE_TYPEINFO(tint::transform::LocalizeStructArrayAssignment);
namespace tint::transform {
/// PIMPL state for the transform
struct LocalizeStructArrayAssignment::State {
/// Constructor
/// @param program the source program
explicit State(const Program* program) : src(program) {}
/// Runs the transform
/// @returns the new program or SkipTransform if the transform is not required
ApplyResult Run() {
struct Shared {
bool process_nested_nodes = false;
utils::Vector<const ast::Statement*, 4> insert_before_stmts;
utils::Vector<const ast::Statement*, 4> insert_after_stmts;
} s;
ctx.ReplaceAll([&](const ast::AssignmentStatement* assign_stmt) -> const ast::Statement* {
// Process if it's an assignment statement to a dynamically indexed array
// within a struct on a function or private storage variable. This
// specific use-case is what FXC fails to compile with:
// error X3500: array reference cannot be used as an l-value; not natively
// addressable
if (!ContainsStructArrayIndex(assign_stmt->lhs)) {
return nullptr;
}
auto og = GetOriginatingTypeAndAddressSpace(assign_stmt);
if (!(og.first->Is<sem::Struct>() && (og.second == ast::AddressSpace::kFunction ||
og.second == ast::AddressSpace::kPrivate))) {
return nullptr;
}
// Reset shared state for this assignment statement
s = Shared{};
const ast::Expression* new_lhs = nullptr;
{
TINT_SCOPED_ASSIGNMENT(s.process_nested_nodes, true);
new_lhs = ctx.Clone(assign_stmt->lhs);
}
auto* new_assign_stmt = b.Assign(new_lhs, ctx.Clone(assign_stmt->rhs));
// Combine insert_before_stmts + new_assign_stmt + insert_after_stmts into
// a block and return it
auto stmts = std::move(s.insert_before_stmts);
stmts.Reserve(1 + s.insert_after_stmts.Length());
stmts.Push(new_assign_stmt);
for (auto* stmt : s.insert_after_stmts) {
stmts.Push(stmt);
}
return b.Block(std::move(stmts));
});
ctx.ReplaceAll(
[&](const ast::IndexAccessorExpression* index_access) -> const ast::Expression* {
if (!s.process_nested_nodes) {
return nullptr;
}
// Indexing a member access expr?
auto* mem_access = index_access->object->As<ast::MemberAccessorExpression>();
if (!mem_access) {
return nullptr;
}
// Process any nested IndexAccessorExpressions
mem_access = ctx.Clone(mem_access);
// Store the address of the member access into a let as we need to read
// the value twice e.g. let tint_symbol = &(s.a1);
auto mem_access_ptr = b.Sym();
s.insert_before_stmts.Push(b.Decl(b.Let(mem_access_ptr, b.AddressOf(mem_access))));
// Disable further transforms when cloning
TINT_SCOPED_ASSIGNMENT(s.process_nested_nodes, false);
// Copy entire array out of struct into local temp var
// e.g. var tint_symbol_1 = *(tint_symbol);
auto tmp_var = b.Sym();
s.insert_before_stmts.Push(b.Decl(b.Var(tmp_var, b.Deref(mem_access_ptr))));
// Replace input index_access with a clone of itself, but with its
// .object replaced by the new temp var. This is returned from this
// function to modify the original assignment statement. e.g.
// tint_symbol_1[uniforms.i]
auto* new_index_access = b.IndexAccessor(tmp_var, ctx.Clone(index_access->index));
// Assign temp var back to array
// e.g. *(tint_symbol) = tint_symbol_1;
auto* assign_rhs_to_temp = b.Assign(b.Deref(mem_access_ptr), tmp_var);
{
utils::Vector<const ast::Statement*, 8> stmts{assign_rhs_to_temp};
for (auto* stmt : s.insert_after_stmts) {
stmts.Push(stmt);
}
s.insert_after_stmts = std::move(stmts);
}
return new_index_access;
});
ctx.Clone();
return Program(std::move(b));
}
private:
/// The source program
const Program* const src;
/// The target program builder
ProgramBuilder b;
/// The clone context
CloneContext ctx = {&b, src, /* auto_clone_symbols */ true};
/// Returns true if `expr` contains an index accessor expression to a
/// structure member of array type.
bool ContainsStructArrayIndex(const ast::Expression* expr) {
bool result = false;
ast::TraverseExpressions(
expr, b.Diagnostics(), [&](const ast::IndexAccessorExpression* ia) {
// Indexing using a runtime value?
auto* idx_sem = src->Sem().Get(ia->index);
if (!idx_sem->ConstantValue()) {
// Indexing a member access expr?
if (auto* ma = ia->object->As<ast::MemberAccessorExpression>()) {
// That accesses an array?
if (src->TypeOf(ma)->UnwrapRef()->Is<sem::Array>()) {
result = true;
return ast::TraverseAction::Stop;
}
}
}
return ast::TraverseAction::Descend;
});
return result;
}
// Returns the type and address space of the originating variable of the lhs
// of the assignment statement.
// See https://www.w3.org/TR/WGSL/#originating-variable-section
std::pair<const sem::Type*, ast::AddressSpace> GetOriginatingTypeAndAddressSpace(
const ast::AssignmentStatement* assign_stmt) {
auto* source_var = src->Sem().Get(assign_stmt->lhs)->SourceVariable();
if (!source_var) {
TINT_ICE(Transform, b.Diagnostics())
<< "Unable to determine originating variable for lhs of assignment "
"statement";
return {};
}
auto* type = source_var->Type();
if (auto* ref = type->As<sem::Reference>()) {
return {ref->StoreType(), ref->AddressSpace()};
} else if (auto* ptr = type->As<sem::Pointer>()) {
return {ptr->StoreType(), ptr->AddressSpace()};
}
TINT_ICE(Transform, b.Diagnostics())
<< "Expecting to find variable of type pointer or reference on lhs "
"of assignment statement";
return {};
}
};
LocalizeStructArrayAssignment::LocalizeStructArrayAssignment() = default;
LocalizeStructArrayAssignment::~LocalizeStructArrayAssignment() = default;
Transform::ApplyResult LocalizeStructArrayAssignment::Apply(const Program* src,
const DataMap&,
DataMap&) const {
return State{src}.Run();
}
} // namespace tint::transform