HLSL: work around FXC failures when dynamically indexing arrays in structs

FXC fails to compile code that assigns to dynamically-indexed fixed-size
arrays in structs on internal shader variables with:

error X3500: array reference cannot be used as an l-value; not natively
addressable

This CL detects this case, and transforms such assignments into copying
out the array to a local variable, assigning to that local, and then
copying the array back.

Also manually regenerate SKIPs for HLSL/FXC after this change, which
fixes 30 tests. Also exposes some "compilation aborted unexpectedly" now
that  "array reference cannot be used as an l-value" has been fixed. For
tests that fail for both DXC and FXC, updating SKIPs to the DXC one to
help distinguish actual FXC bugs from valid errors.

Bug: tint:998
Bug: tint:1206
Change-Id: I09204d8d81ab27d1c257538ad702414ccc386543
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/71620
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>

src/BUILD.gn

@@ -451,6 +451,8 @@ libtint_source_set("libtint_core_all_src") {
     "transform/fold_trivial_single_use_lets.h",
     "transform/for_loop_to_loop.cc",
     "transform/for_loop_to_loop.h",
+    "transform/localize_struct_array_assignment.cc",
+    "transform/localize_struct_array_assignment.h",
     "transform/loop_to_for_loop.cc",
     "transform/loop_to_for_loop.h",
     "transform/manager.cc",

src/CMakeLists.txt

@@ -312,6 +312,8 @@ set(TINT_LIB_SRCS
   transform/fold_constants.h
   transform/fold_trivial_single_use_lets.cc
   transform/fold_trivial_single_use_lets.h
+  transform/localize_struct_array_assignment.cc
+  transform/localize_struct_array_assignment.h
   transform/for_loop_to_loop.cc
   transform/for_loop_to_loop.h
   transform/glsl.cc
@@ -971,6 +973,7 @@ if(${TINT_BUILD_TESTS})
       transform/fold_constants_test.cc
       transform/fold_trivial_single_use_lets_test.cc
       transform/for_loop_to_loop_test.cc
+      transform/localize_struct_array_assignment_test.cc
       transform/loop_to_for_loop_test.cc
       transform/module_scope_var_to_entry_point_param_test.cc
       transform/multiplanar_external_texture_test.cc

src/program_builder.h

@@ -2468,6 +2468,15 @@ class ProgramBuilder {
     }
   }
 
+  /// Unmarks that the given transform `T` has been applied to this program.
+  template <typename T>
+  void UnsetTransformApplied() {
+    auto it = transforms_applied_.find(&TypeInfo::Of<T>());
+    if (it != transforms_applied_.end()) {
+      transforms_applied_.erase(it);
+    }
+  }
+
   /// @returns true if the transform of type `T` was applied.
   template <typename T>
   bool HasTransformApplied() {

src/transform/localize_struct_array_assignment.cc

@@ -0,0 +1,230 @@
+// 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/localize_struct_array_assignment.h"
+
+#include <unordered_map>
+#include <utility>
+
+#include "src/ast/assignment_statement.h"
+#include "src/ast/traverse_expressions.h"
+#include "src/program_builder.h"
+#include "src/sem/expression.h"
+#include "src/sem/member_accessor_expression.h"
+#include "src/sem/reference_type.h"
+#include "src/sem/statement.h"
+#include "src/sem/variable.h"
+#include "src/transform/simplify_pointers.h"
+#include "src/utils/scoped_assignment.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::transform::LocalizeStructArrayAssignment);
+
+namespace tint {
+namespace transform {
+
+/// Private implementation of LocalizeStructArrayAssignment transform
+class LocalizeStructArrayAssignment::State {
+ private:
+  CloneContext& ctx;
+  ProgramBuilder& b;
+
+  /// 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 = ctx.src->Sem().Get(ia->index);
+          if (!idx_sem->ConstantValue().IsValid()) {
+            // Indexing a member access expr?
+            if (auto* ma = ia->object->As<ast::MemberAccessorExpression>()) {
+              // That accesses an array?
+              if (ctx.src->TypeOf(ma)->UnwrapRef()->Is<sem::Array>()) {
+                result = true;
+                return ast::TraverseAction::Stop;
+              }
+            }
+          }
+          return ast::TraverseAction::Descend;
+        });
+
+    return result;
+  }
+
+  // Returns the type and storage class 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::StorageClass>
+  GetOriginatingTypeAndStorageClass(
+      const ast::AssignmentStatement* assign_stmt) {
+    // Get first IdentifierExpr from lhs of assignment, which should resolve to
+    // the pointer or reference of the originating variable of the assignment.
+    // TraverseExpressions traverses left to right, and this code depends on the
+    // fact that for an assignment statement, the variable will be the left-most
+    // expression.
+    // TODO(crbug.com/tint/1341): do this in the Resolver, setting the
+    // originating variable on sem::Expression.
+    const ast::IdentifierExpression* ident = nullptr;
+    ast::TraverseExpressions(assign_stmt->lhs, b.Diagnostics(),
+                             [&](const ast::IdentifierExpression* id) {
+                               ident = id;
+                               return ast::TraverseAction::Stop;
+                             });
+    auto* sem_var_user = ctx.src->Sem().Get<sem::VariableUser>(ident);
+    if (!sem_var_user) {
+      TINT_ICE(Transform, b.Diagnostics())
+          << "Expected to find variable of lhs of assignment statement";
+      return {};
+    }
+
+    auto* var = sem_var_user->Variable();
+    if (auto* ptr = var->Type()->As<sem::Pointer>()) {
+      return {ptr->StoreType(), ptr->StorageClass()};
+    }
+
+    auto* ref = var->Type()->As<sem::Reference>();
+    if (!ref) {
+      TINT_ICE(Transform, b.Diagnostics())
+          << "Expecting to find variable of type pointer or reference on lhs "
+             "of assignment statement";
+      return {};
+    }
+
+    return {ref->StoreType(), ref->StorageClass()};
+  }
+
+ public:
+  /// Constructor
+  /// @param ctx_in the CloneContext primed with the input program and
+  /// ProgramBuilder
+  explicit State(CloneContext& ctx_in) : ctx(ctx_in), b(*ctx_in.dst) {}
+
+  /// Runs the transform
+  void Run() {
+    struct Shared {
+      bool process_nested_nodes = false;
+      ast::StatementList insert_before_stmts;
+      ast::StatementList 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 = GetOriginatingTypeAndStorageClass(assign_stmt);
+      if (!(og.first->Is<sem::Struct>() &&
+            (og.second == ast::StorageClass::kFunction ||
+             og.second == ast::StorageClass::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
+      ast::StatementList stmts = std::move(s.insert_before_stmts);
+      stmts.reserve(1 + s.insert_after_stmts.size());
+      stmts.emplace_back(new_assign_stmt);
+      stmts.insert(stmts.end(), s.insert_after_stmts.begin(),
+                   s.insert_after_stmts.end());
+
+      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_back(
+          b.Decl(b.Const(mem_access_ptr, nullptr, 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_back(
+          b.Decl(b.Var(tmp_var, nullptr, 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);
+      s.insert_after_stmts.insert(s.insert_after_stmts.begin(),
+                                  assign_rhs_to_temp);  // push_front
+
+      return new_index_access;
+    });
+
+    ctx.Clone();
+  }
+};
+
+LocalizeStructArrayAssignment::LocalizeStructArrayAssignment() = default;
+
+LocalizeStructArrayAssignment::~LocalizeStructArrayAssignment() = default;
+
+void LocalizeStructArrayAssignment::Run(CloneContext& ctx,
+                                        const DataMap&,
+                                        DataMap&) {
+  if (!Requires<SimplifyPointers>(ctx)) {
+    return;
+  }
+
+  State state(ctx);
+  state.Run();
+
+  // This transform may introduce pointers
+  ctx.dst->UnsetTransformApplied<transform::SimplifyPointers>();
+}
+}  // namespace transform
+}  // namespace tint

src/transform/localize_struct_array_assignment.h

@@ -0,0 +1,53 @@
+// 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_LOCALIZE_STRUCT_ARRAY_ASSIGNMENT_H_
+#define SRC_TRANSFORM_LOCALIZE_STRUCT_ARRAY_ASSIGNMENT_H_
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// This transforms replaces assignment to dynamically-indexed fixed-size arrays
+/// in structs on shader-local variables with code that copies the arrays to a
+/// temporary local variable, assigns to the local variable, and copies the
+/// array back. This is to work around FXC's compilation failure for these cases
+/// (see crbug.com/tint/1206).
+class LocalizeStructArrayAssignment
+    : public Castable<LocalizeStructArrayAssignment, Transform> {
+ public:
+  /// Constructor
+  LocalizeStructArrayAssignment();
+
+  /// Destructor
+  ~LocalizeStructArrayAssignment() override;
+
+ protected:
+  /// Runs the transform using the CloneContext built for transforming a
+  /// program. Run() is responsible for calling Clone() on the CloneContext.
+  /// @param ctx the CloneContext primed with the input program and
+  /// ProgramBuilder
+  /// @param inputs optional extra transform-specific input data
+  /// @param outputs optional extra transform-specific output data
+  void Run(CloneContext& ctx, const DataMap& inputs, DataMap& outputs) override;
+
+ private:
+  class State;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_LOCALIZE_STRUCT_ARRAY_ASSIGNMENT_H_

src/transform/localize_struct_array_assignment_test.cc

@@ -0,0 +1,620 @@
+// 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/localize_struct_array_assignment.h"
+#include "src/transform/simplify_pointers.h"
+#include "src/transform/unshadow.h"
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using LocalizeStructArrayAssignmentTest = TransformTest;
+
+TEST_F(LocalizeStructArrayAssignmentTest, MissingSimplifyPointers) {
+  auto* src = R"()";
+  auto* expect =
+      "error: tint::transform::LocalizeStructArrayAssignment depends on "
+      "tint::transform::SimplifyPointers but the dependency was not run";
+
+  auto got = Run<LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, EmptyModule) {
+  auto* src = R"()";
+  auto* expect = src;
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, StructArray) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  s1.a1[uniforms.i] = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  {
+    let tint_symbol = &(s1.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[uniforms.i] = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, StructStructArray) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a : array<InnerS, 8>;
+};
+
+struct OuterS {
+  s2 : S1;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  s1.s2.a[uniforms.i] = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a : array<InnerS, 8>;
+};
+
+struct OuterS {
+  s2 : S1;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  {
+    let tint_symbol = &(s1.s2.a);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[uniforms.i] = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, StructArrayArray) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<array<InnerS, 8>, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  s1.a1[uniforms.i][uniforms.j] = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<array<InnerS, 8>, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  {
+    let tint_symbol = &(s1.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[uniforms.i][uniforms.j] = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, StructArrayStruct) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  s2 : InnerS;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  s1.a1[uniforms.i].s2 = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  s2 : InnerS;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  {
+    let tint_symbol = &(s1.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[uniforms.i].s2 = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, StructArrayStructArray) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a2 : array<InnerS, 8>;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s : OuterS;
+  s.a1[uniforms.i].a2[uniforms.j] = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a2 : array<InnerS, 8>;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s : OuterS;
+  {
+    let tint_symbol = &(s.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    let tint_symbol_2 = &(tint_symbol_1[uniforms.i].a2);
+    var tint_symbol_3 = *(tint_symbol_2);
+    tint_symbol_3[uniforms.j] = v;
+    *(tint_symbol_2) = tint_symbol_3;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, IndexingWithSideEffectFunc) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a2 : array<InnerS, 8>;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+var<private> nextIndex : u32;
+fn getNextIndex() -> u32 {
+  nextIndex = nextIndex + 1u;
+  return nextIndex;
+}
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s : OuterS;
+  s.a1[getNextIndex()].a2[uniforms.j] = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+  j : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct S1 {
+  a2 : array<InnerS, 8>;
+};
+
+struct OuterS {
+  a1 : array<S1, 8>;
+};
+
+var<private> nextIndex : u32;
+
+fn getNextIndex() -> u32 {
+  nextIndex = (nextIndex + 1u);
+  return nextIndex;
+}
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s : OuterS;
+  {
+    let tint_symbol = &(s.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    let tint_symbol_2 = &(tint_symbol_1[getNextIndex()].a2);
+    var tint_symbol_3 = *(tint_symbol_2);
+    tint_symbol_3[uniforms.j] = v;
+    *(tint_symbol_2) = tint_symbol_3;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, ViaPointerArg) {
+  auto* src = R"(
+[[block]] struct Uniforms {
+  i : u32;
+};
+struct InnerS {
+  v : i32;
+};
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+fn f(p : ptr<function, OuterS>) {
+  var v : InnerS;
+  (*p).a1[uniforms.i] = v;
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var s1 : OuterS;
+  f(&s1);
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+fn f(p : ptr<function, OuterS>) {
+  var v : InnerS;
+  {
+    let tint_symbol = &((*(p)).a1);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[uniforms.i] = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var s1 : OuterS;
+  f(&(s1));
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, ViaPointerVar) {
+  auto* src = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+fn f(p : ptr<function, InnerS>, v : InnerS) {
+  *(p) = v;
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  let p = &(s1.a1[uniforms.i]);
+  *(p) = v;
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+struct InnerS {
+  v : i32;
+};
+
+struct OuterS {
+  a1 : array<InnerS, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+fn f(p : ptr<function, InnerS>, v : InnerS) {
+  *(p) = v;
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var v : InnerS;
+  var s1 : OuterS;
+  let p_save = uniforms.i;
+  {
+    let tint_symbol = &(s1.a1);
+    var tint_symbol_1 = *(tint_symbol);
+    tint_symbol_1[p_save] = v;
+    *(tint_symbol) = tint_symbol_1;
+  }
+}
+)";
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(LocalizeStructArrayAssignmentTest, VectorAssignment) {
+  auto* src = R"(
+[[block]]
+struct Uniforms {
+  i : u32;
+};
+
+[[block]]
+struct OuterS {
+  a1 : array<u32, 8>;
+};
+
+[[group(1), binding(4)]] var<uniform> uniforms : Uniforms;
+
+fn f(i : u32) -> u32 {
+  return (i + 1u);
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+  var s1 : OuterS;
+  var v : vec3<f32>;
+  v[s1.a1[uniforms.i]] = 1.0;
+  v[f(s1.a1[uniforms.i])] = 1.0;
+}
+)";
+
+  // Transform does nothing here as we're not actually assigning to the array in
+  // the struct.
+  auto* expect = src;
+
+  auto got =
+      Run<Unshadow, SimplifyPointers, LocalizeStructArrayAssignment>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint

src/writer/hlsl/generator_impl.cc

@@ -51,6 +51,7 @@
 #include "src/transform/decompose_memory_access.h"
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/fold_trivial_single_use_lets.h"
+#include "src/transform/localize_struct_array_assignment.h"
 #include "src/transform/loop_to_for_loop.h"
 #include "src/transform/manager.h"
 #include "src/transform/num_workgroups_from_uniform.h"
@@ -145,6 +146,15 @@ SanitizedResult Sanitize(
 
   manager.Add<transform::Unshadow>();
 
+  // LocalizeStructArrayAssignment must come after:
+  // * SimplifyPointers, because it assumes assignment to arrays in structs are
+  // done directly, not indirectly.
+  // TODO(crbug.com/tint/1340): See if we can get rid of the duplicate
+  // SimplifyPointers transform. Can't do it right now because
+  // LocalizeStructArrayAssignment introduces pointers.
+  manager.Add<transform::SimplifyPointers>();
+  manager.Add<transform::LocalizeStructArrayAssignment>();
+
   // Attempt to convert `loop`s into for-loops. This is to try and massage the
   // output into something that will not cause FXC to choke or misbehave.
   manager.Add<transform::FoldTrivialSingleUseLets>();