Add transform/WrapArraysInStructs

And replace the MSL writer's logic to do this with the transform.

We need to do the same thing in HLSL, and in the future GLSL too.

Partially reverts fbfde720

Change-Id: Ie280e011bc3ded8e15ccacc0aeb12da3c2407389
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/54242
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/transform/msl.cc

@@ -28,6 +28,7 @@
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/manager.h"
 #include "src/transform/promote_initializers_to_const_var.h"
+#include "src/transform/wrap_arrays_in_structs.h"
 
 namespace tint {
 namespace transform {
@@ -41,6 +42,7 @@ Output Msl::Run(const Program* in, const DataMap&) {
   manager.Add<CanonicalizeEntryPointIO>();
   manager.Add<ExternalTextureTransform>();
   manager.Add<PromoteInitializersToConstVar>();
+  manager.Add<WrapArraysInStructs>();
   data.Add<CanonicalizeEntryPointIO::Config>(
       CanonicalizeEntryPointIO::BuiltinStyle::kParameter);
   auto out = manager.Run(in, data);

src/transform/wrap_arrays_in_structs.cc

@@ -0,0 +1,152 @@
+// 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

src/transform/wrap_arrays_in_structs.h

@@ -0,0 +1,84 @@
+// 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.
+
+#ifndef SRC_TRANSFORM_WRAP_ARRAYS_IN_STRUCTS_H_
+#define SRC_TRANSFORM_WRAP_ARRAYS_IN_STRUCTS_H_
+
+#include <string>
+#include <unordered_map>
+
+#include "src/transform/transform.h"
+
+// Forward declarations
+namespace tint {
+namespace ast {
+class Type;
+}  // namespace ast
+}  // namespace tint
+
+namespace tint {
+namespace transform {
+
+/// WrapArraysInStructs is a transform that replaces all array types with a
+/// structure holding a single field of that array type.
+/// Array index expressions and constructors are also adjusted to deal with this
+/// wrapping.
+/// This transform helps with backends that cannot directly return arrays or use
+/// them as parameters.
+class WrapArraysInStructs : public Transform {
+ public:
+  /// Constructor
+  WrapArraysInStructs();
+
+  /// Destructor
+  ~WrapArraysInStructs() override;
+
+  /// Runs the transform on `program`, returning the transformation result.
+  /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
+  /// @returns the transformation result
+  Output Run(const Program* program, const DataMap& data = {}) override;
+
+ private:
+  struct WrappedArrayInfo {
+    WrappedArrayInfo();
+    WrappedArrayInfo(const WrappedArrayInfo&);
+    ~WrappedArrayInfo();
+
+    Symbol wrapper_name;
+    std::function<ast::Type*(CloneContext&)> array_type;
+
+    operator bool() { return wrapper_name.IsValid(); }
+  };
+
+  /// WrapArray wraps the fixed-size array type in a new structure (if it hasn't
+  /// already been wrapped). WrapArray will recursively wrap arrays-of-arrays.
+  /// The new structure will be added to module-scope type declarations of
+  /// `ctx.dst`.
+  /// @param ctx the CloneContext
+  /// @param wrapped_arrays a map of src array type to the wrapped structure
+  /// name
+  /// @param array the array type
+  /// @return the name of the structure that wraps the array, or an invalid
+  /// Symbol if this array should not be wrapped
+  WrappedArrayInfo WrapArray(
+      CloneContext& ctx,
+      std::unordered_map<const sem::Array*, WrappedArrayInfo>& wrapped_arrays,
+      const sem::Array* array) const;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_WRAP_ARRAYS_IN_STRUCTS_H_

src/transform/wrap_arrays_in_structs_test.cc

@@ -0,0 +1,317 @@
+// 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 <memory>
+#include <utility>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using WrapArraysInStructsTest = TransformTest;
+
+TEST_F(WrapArraysInStructsTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArrayAsGlobal) {
+  auto* src = R"(
+var<private> arr : array<i32, 4>;
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+var<private> arr : tint_array_wrapper;
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArrayAsFunctionVar) {
+  auto* src = R"(
+fn f() {
+  var arr : array<i32, 4>;
+  let x = arr[3];
+}
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+fn f() {
+  var arr : tint_array_wrapper;
+  let x = arr.arr[3];
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArrayAsParam) {
+  auto* src = R"(
+fn f(a : array<i32, 4>) -> i32 {
+  return a[2];
+}
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+fn f(a : tint_array_wrapper) -> i32 {
+  return a.arr[2];
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArrayAsReturn) {
+  auto* src = R"(
+fn f() -> array<i32, 4> {
+  return array<i32, 4>(1, 2, 3, 4);
+}
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+fn f() -> tint_array_wrapper {
+  return tint_array_wrapper(array<i32, 4>(1, 2, 3, 4));
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArrayAlias) {
+  auto* src = R"(
+type Inner = array<i32, 2>;
+type Array = array<Inner, 2>;
+
+fn f() {
+  var arr : Array;
+  arr = Array();
+  arr = Array(Inner(1, 2), Inner(3, 4));
+  let vals : Array = Array(Inner(1, 2), Inner(3, 4));
+  arr = vals;
+  let x = arr[3];
+}
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 2>;
+};
+
+type Inner = tint_array_wrapper;
+
+struct tint_array_wrapper_1 {
+  arr : array<tint_array_wrapper, 2>;
+};
+
+type Array = tint_array_wrapper_1;
+
+fn f() {
+  var arr : tint_array_wrapper_1;
+  arr = tint_array_wrapper_1(array<tint_array_wrapper, 2>());
+  arr = tint_array_wrapper_1(array<tint_array_wrapper, 2>(tint_array_wrapper(array<i32, 2>(1, 2)), tint_array_wrapper(array<i32, 2>(3, 4))));
+  let vals : tint_array_wrapper_1 = tint_array_wrapper_1(array<tint_array_wrapper, 2>(tint_array_wrapper(array<i32, 2>(1, 2)), tint_array_wrapper(array<i32, 2>(3, 4))));
+  arr = vals;
+  let x = arr.arr[3];
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : array<i32, 4>;
+  b : array<i32, 8>;
+  c : array<i32, 4>;
+};
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+struct tint_array_wrapper_1 {
+  arr : array<i32, 8>;
+};
+
+struct S {
+  a : tint_array_wrapper;
+  b : tint_array_wrapper_1;
+  c : tint_array_wrapper;
+};
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, ArraysOfArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : array<i32, 4>;
+  b : array<array<i32, 4>, 4>;
+  c : array<array<array<i32, 4>, 4>, 4>;
+};
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+struct tint_array_wrapper_1 {
+  arr : array<tint_array_wrapper, 4>;
+};
+
+struct tint_array_wrapper_2 {
+  arr : array<tint_array_wrapper_1, 4>;
+};
+
+struct S {
+  a : tint_array_wrapper;
+  b : tint_array_wrapper_1;
+  c : tint_array_wrapper_2;
+};
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, AccessArraysOfArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : array<i32, 4>;
+  b : array<array<i32, 4>, 4>;
+  c : array<array<array<i32, 4>, 4>, 4>;
+};
+
+fn f(s : S) -> i32 {
+  return s.a[2] + s.b[1][2] + s.c[3][1][2];
+}
+)";
+  auto* expect = R"(
+struct tint_array_wrapper {
+  arr : array<i32, 4>;
+};
+
+struct tint_array_wrapper_1 {
+  arr : array<tint_array_wrapper, 4>;
+};
+
+struct tint_array_wrapper_2 {
+  arr : array<tint_array_wrapper_1, 4>;
+};
+
+struct S {
+  a : tint_array_wrapper;
+  b : tint_array_wrapper_1;
+  c : tint_array_wrapper_2;
+};
+
+fn f(s : S) -> i32 {
+  return ((s.a.arr[2] + s.b.arr[1].arr[2]) + s.c.arr[3].arr[1].arr[2]);
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(WrapArraysInStructsTest, DeclarationOrder) {
+  auto* src = R"(
+type T0 = i32;
+
+type T1 = array<i32, 1>;
+
+type T2 = i32;
+
+fn f1(a : array<i32, 2>) {
+}
+
+type T3 = i32;
+
+fn f2() {
+  var v : array<i32, 3>;
+}
+)";
+  auto* expect = R"(
+type T0 = i32;
+
+struct tint_array_wrapper {
+  arr : array<i32, 1>;
+};
+
+type T1 = tint_array_wrapper;
+
+type T2 = i32;
+
+struct tint_array_wrapper_1 {
+  arr : array<i32, 2>;
+};
+
+fn f1(a : tint_array_wrapper_1) {
+}
+
+type T3 = i32;
+
+struct tint_array_wrapper_2 {
+  arr : array<i32, 3>;
+};
+
+fn f2() {
+  var v : tint_array_wrapper_2;
+}
+)";
+
+  auto got = Run<WrapArraysInStructs>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint