blob: b1715aa15fa828d1d9393287bf58365e60fe5bf7 [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/transform/wrap_arrays_in_structs.h"
#include <utility>
#include "src/program_builder.h"
#include "src/sem/array.h"
#include "src/sem/expression.h"
#include "src/utils/get_or_create.h"
namespace tint {
namespace transform {
WrapArraysInStructs::WrappedArrayInfo::WrappedArrayInfo() = default;
WrapArraysInStructs::WrappedArrayInfo::WrappedArrayInfo(
const WrappedArrayInfo&) = default;
WrapArraysInStructs::WrappedArrayInfo::~WrappedArrayInfo() = default;
WrapArraysInStructs::WrapArraysInStructs() = default;
WrapArraysInStructs::~WrapArraysInStructs() = default;
Output WrapArraysInStructs::Run(const Program* in, const DataMap&) {
ProgramBuilder out;
CloneContext ctx(&out, in);
auto& sem = ctx.src->Sem();
std::unordered_map<const sem::Array*, WrappedArrayInfo> wrapped_arrays;
auto wrapper = [&](const sem::Array* array) {
return WrapArray(ctx, wrapped_arrays, array);
};
auto wrapper_typename = [&](const sem::Array* arr) -> ast::TypeName* {
auto info = wrapper(arr);
return info ? ctx.dst->create<ast::TypeName>(info.wrapper_name) : nullptr;
};
// Replace all array types with their corresponding wrapper
ctx.ReplaceAll([&](ast::Type* ast_type) -> ast::Type* {
auto* type = ctx.src->TypeOf(ast_type);
if (auto* array = type->UnwrapRef()->As<sem::Array>()) {
return wrapper_typename(array);
}
return nullptr;
});
// Fix up array accessors so `a[1]` becomes `a.arr[1]`
ctx.ReplaceAll([&](ast::ArrayAccessorExpression* accessor)
-> ast::ArrayAccessorExpression* {
if (auto* array =
As<sem::Array>(sem.Get(accessor->array())->Type()->UnwrapRef())) {
if (wrapper(array)) {
// Array is wrapped in a structure. Emit a member accessor to get
// to the actual array.
auto* arr = ctx.Clone(accessor->array());
auto* idx = ctx.Clone(accessor->idx_expr());
auto* unwrapped = ctx.dst->MemberAccessor(arr, "arr");
return ctx.dst->IndexAccessor(accessor->source(), unwrapped, idx);
}
}
return nullptr;
});
// Fix up array constructors so `A(1,2)` becomes `tint_array_wrapper(A(1,2))`
ctx.ReplaceAll([&](ast::TypeConstructorExpression* ctor) -> ast::Expression* {
if (auto* array = As<sem::Array>(sem.Get(ctor)->Type()->UnwrapRef())) {
if (auto w = wrapper(array)) {
// Wrap the array type constructor with another constructor for
// the wrapper
auto* wrapped_array_ty = ctx.Clone(ctor->type());
auto* array_ty = w.array_type(ctx);
auto* arr_ctor =
ctx.dst->Construct(array_ty, ctx.Clone(ctor->values()));
return ctx.dst->Construct(wrapped_array_ty, arr_ctor);
}
}
return nullptr;
});
ctx.Clone();
return Output(Program(std::move(out)));
}
WrapArraysInStructs::WrappedArrayInfo WrapArraysInStructs::WrapArray(
CloneContext& ctx,
std::unordered_map<const sem::Array*, WrappedArrayInfo>& wrapped_arrays,
const sem::Array* array) const {
if (array->IsRuntimeSized()) {
return {}; // We don't want to wrap runtime sized arrays
}
return utils::GetOrCreate(wrapped_arrays, array, [&] {
WrappedArrayInfo info;
// Generate a unique name for the array wrapper
info.wrapper_name = ctx.dst->Symbols().New("tint_array_wrapper");
// Examine the element type. Is it also an array?
std::function<ast::Type*(CloneContext&)> el_type;
if (auto* el_array = array->ElemType()->As<sem::Array>()) {
// Array of array - call WrapArray() on the element type
if (auto el = WrapArray(ctx, wrapped_arrays, el_array)) {
el_type = [=](CloneContext& c) {
return c.dst->create<ast::TypeName>(el.wrapper_name);
};
}
}
// If the element wasn't an array, just create the typical AST type for it
if (!el_type) {
el_type = [=](CloneContext& c) {
return CreateASTTypeFor(&c, array->ElemType());
};
}
// Construct the single structure field type
info.array_type = [=](CloneContext& c) {
ast::DecorationList decos;
if (!array->IsStrideImplicit()) {
decos.emplace_back(
c.dst->create<ast::StrideDecoration>(array->Stride()));
}
return c.dst->create<ast::Array>(el_type(c), array->Count(),
std::move(decos));
};
// Structure() will create and append the ast::Struct to the
// global declarations of `ctx.dst`. As we haven't finished building the
// current module-scope statement or function, this will be placed
// immediately before the usage.
ctx.dst->Structure(info.wrapper_name,
{ctx.dst->Member("arr", info.array_type(ctx))});
return info;
});
}
} // namespace transform
} // namespace tint