writer/wgsl: Fix size / align decoration emission

This was broken by a rebase of the Default Struct Layout change.
This went unnoticed because there was no test coverage for these. Added.

Also replace `[[offset(n)]]` decorations with padding fields.

Bug: tint:626
Change-Id: Iad6f1a239bc8d8fcb15d18a204d3f5a78a372350
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44683
Commit-Queue: Ben Clayton <bclayton@chromium.org>
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/resolver/resolver.cc

@@ -1188,7 +1188,7 @@ const semantic::Struct* Resolver::Structure(type::Struct* str) {
     offset = utils::RoundUp(align, offset);
 
     auto* sem_member =
-        builder_->create<semantic::StructMember>(member, offset, size);
+        builder_->create<semantic::StructMember>(member, offset, align, size);
     builder_->Sem().Add(member, sem_member);
     sem_members.emplace_back(sem_member);
 

src/semantic/sem_struct.cc

@@ -30,8 +30,9 @@ Struct::~Struct() = default;
 
 StructMember::StructMember(ast::StructMember* declaration,
                            uint32_t offset,
+                           uint32_t align,
                            uint32_t size)
-    : declaration_(declaration), offset_(offset), size_(size) {}
+    : declaration_(declaration), offset_(offset), align_(align), size_(size) {}
 
 StructMember::~StructMember() = default;
 

src/semantic/struct.h

@@ -85,8 +85,12 @@ class StructMember : public Castable<StructMember, Node> {
   /// Constructor
   /// @param declaration the AST declaration node
   /// @param offset the byte offset from the base of the structure
-  /// @param size the byte size
-  StructMember(ast::StructMember* declaration, uint32_t offset, uint32_t size);
+  /// @param align the byte alignment of the member
+  /// @param size the byte size of the member
+  StructMember(ast::StructMember* declaration,
+               uint32_t offset,
+               uint32_t align,
+               uint32_t size);
 
   /// Destructor
   ~StructMember() override;
@@ -97,13 +101,17 @@ class StructMember : public Castable<StructMember, Node> {
   /// @returns byte offset from base of structure
   uint32_t Offset() const { return offset_; }
 
+  /// @returns the alignment of the member in bytes
+  uint32_t Align() const { return align_; }
+
   /// @returns byte size
   uint32_t Size() const { return size_; }
 
  private:
   ast::StructMember* const declaration_;
   uint32_t const offset_;  // Byte offset from base of structure
-  uint32_t const size_;    // Byte size
+  uint32_t const align_;   // Byte alignment of the member
+  uint32_t const size_;    // Byte size of the member
 };
 
 }  // namespace semantic

src/writer/wgsl/generator_impl.cc

@@ -15,6 +15,7 @@
 #include "src/writer/wgsl/generator_impl.h"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/ast/bool_literal.h"
 #include "src/ast/call_statement.h"
@@ -31,6 +32,7 @@
 #include "src/ast/variable_decl_statement.h"
 #include "src/ast/workgroup_decoration.h"
 #include "src/semantic/function.h"
+#include "src/semantic/struct.h"
 #include "src/semantic/variable.h"
 #include "src/type/access_control_type.h"
 #include "src/type/alias_type.h"
@@ -46,6 +48,7 @@
 #include "src/type/u32_type.h"
 #include "src/type/vector_type.h"
 #include "src/type/void_type.h"
+#include "src/utils/math.h"
 #include "src/writer/float_to_string.h"
 
 namespace tint {
@@ -503,11 +506,41 @@ bool GeneratorImpl::EmitStructType(const type::Struct* str) {
   out_ << "struct " << program_->Symbols().NameFor(str->symbol()) << " {"
        << std::endl;
 
+  auto add_padding = [&](uint32_t size) {
+    make_indent();
+    out_ << "[[size(" << size << ")]]" << std::endl;
+    make_indent();
+    // Note: u32 is the smallest primitive we currently support. When WGSL
+    // supports smaller types, this will need to be updated.
+    out_ << UniqueIdentifier("padding") << " : u32;" << std::endl;
+  };
+
   increment_indent();
+  uint32_t offset = 0;
   for (auto* mem : impl->members()) {
-    if (!mem->decorations().empty()) {
+    auto* mem_sem = program_->Sem().Get(mem);
+
+    offset = utils::RoundUp(mem_sem->Align(), offset);
+    if (uint32_t padding = mem_sem->Offset() - offset) {
+      add_padding(padding);
+      offset += padding;
+    }
+    offset += mem_sem->Size();
+
+    // Offset decorations no longer exist in the WGSL spec, but are emitted
+    // by the SPIR-V reader and are consumed by the Resolver(). These should not
+    // be emitted, but instead struct padding fields should be emitted.
+    ast::DecorationList decorations_sanitized;
+    decorations_sanitized.reserve(mem->decorations().size());
+    for (auto* deco : mem->decorations()) {
+      if (!deco->Is<ast::StructMemberOffsetDecoration>()) {
+        decorations_sanitized.emplace_back(deco);
+      }
+    }
+
+    if (!decorations_sanitized.empty()) {
       make_indent();
-      if (!EmitDecorations(mem->decorations())) {
+      if (!EmitDecorations(decorations_sanitized)) {
         return false;
       }
       out_ << std::endl;
@@ -585,14 +618,13 @@ bool GeneratorImpl::EmitDecorations(const ast::DecorationList& decos) {
       out_ << "builtin(" << builtin->value() << ")";
     } else if (auto* constant = deco->As<ast::ConstantIdDecoration>()) {
       out_ << "constant_id(" << constant->value() << ")";
-    } else if (auto* offset = deco->As<ast::StructMemberOffsetDecoration>()) {
-      out_ << "offset(" << offset->offset() << ")";
     } else if (auto* size = deco->As<ast::StructMemberSizeDecoration>()) {
-      out_ << "[[size(" << size->size() << ")]]" << std::endl;
+      out_ << "size(" << size->size() << ")";
     } else if (auto* align = deco->As<ast::StructMemberAlignDecoration>()) {
-      out_ << "[[align(" << align->align() << ")]]" << std::endl;
+      out_ << "align(" << align->align() << ")";
     } else {
-      diagnostics_.add_error("unknown variable decoration");
+      TINT_ICE(diagnostics_)
+          << "Unsupported decoration '" << deco->TypeInfo().name << "'";
       return false;
     }
   }
@@ -952,6 +984,23 @@ bool GeneratorImpl::EmitSwitch(ast::SwitchStatement* stmt) {
   return true;
 }
 
+std::string GeneratorImpl::UniqueIdentifier(const std::string& suffix) {
+  auto const limit =
+      std::numeric_limits<decltype(next_unique_identifier_suffix)>::max();
+  while (next_unique_identifier_suffix < limit) {
+    auto ident = "tint_" + std::to_string(next_unique_identifier_suffix);
+    if (!suffix.empty()) {
+      ident += "_" + suffix;
+    }
+    next_unique_identifier_suffix++;
+    if (!program_->Symbols().Get(ident).IsValid()) {
+      return ident;
+    }
+  }
+  diagnostics_.add_error("Unable to generate a unique WGSL identifier");
+  return "<invalid-ident>";
+}
+
 }  // namespace wgsl
 }  // namespace writer
 }  // namespace tint

src/writer/wgsl/generator_impl.h

@@ -199,7 +199,11 @@ class GeneratorImpl : public TextGenerator {
   bool EmitDecorations(const ast::DecorationList& decos);
 
  private:
+  /// @return a new, unique, valid WGSL identifier with the given suffix.
+  std::string UniqueIdentifier(const std::string& suffix = "");
+
   Program const* const program_;
+  uint32_t next_unique_identifier_suffix = 0;
 };
 
 }  // namespace wgsl

src/writer/wgsl/generator_impl_type_test.cc

@@ -159,33 +159,90 @@ TEST_F(WgslGeneratorImplTest, EmitType_Struct) {
   EXPECT_EQ(gen.result(), "S");
 }
 
-TEST_F(WgslGeneratorImplTest, EmitType_StructDecl) {
+TEST_F(WgslGeneratorImplTest, EmitType_StructOffsetDecl) {
   auto* s = Structure("S", {
-                               Member("a", ty.i32()),
-                               Member("b", ty.f32(), {MemberOffset(4)}),
+                               Member("a", ty.i32(), {MemberOffset(8)}),
+                               Member("b", ty.f32(), {MemberOffset(16)}),
                            });
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitStructType(s)) << gen.error();
   EXPECT_EQ(gen.result(), R"(struct S {
+  [[size(8)]]
+  tint_0_padding : u32;
   a : i32;
-  [[offset(4)]]
+  [[size(4)]]
+  tint_1_padding : u32;
+  b : f32;
+};
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, EmitType_StructOffsetDecl_WithSymbolCollisions) {
+  auto* s =
+      Structure("S", {
+                         Member("tint_0_padding", ty.i32(), {MemberOffset(8)}),
+                         Member("tint_2_padding", ty.f32(), {MemberOffset(16)}),
+                     });
+
+  GeneratorImpl& gen = Build();
+
+  ASSERT_TRUE(gen.EmitStructType(s)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(struct S {
+  [[size(8)]]
+  tint_1_padding : u32;
+  tint_0_padding : i32;
+  [[size(4)]]
+  tint_3_padding : u32;
+  tint_2_padding : f32;
+};
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, EmitType_StructAlignDecl) {
+  auto* s = Structure("S", {
+                               Member("a", ty.i32(), {MemberAlign(8)}),
+                               Member("b", ty.f32(), {MemberAlign(16)}),
+                           });
+
+  GeneratorImpl& gen = Build();
+
+  ASSERT_TRUE(gen.EmitStructType(s)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(struct S {
+  [[align(8)]]
+  a : i32;
+  [[align(16)]]
+  b : f32;
+};
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, EmitType_StructSizeDecl) {
+  auto* s = Structure("S", {
+                               Member("a", ty.i32(), {MemberSize(16)}),
+                               Member("b", ty.f32(), {MemberSize(32)}),
+                           });
+
+  GeneratorImpl& gen = Build();
+
+  ASSERT_TRUE(gen.EmitStructType(s)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(struct S {
+  [[size(16)]]
+  a : i32;
+  [[size(32)]]
   b : f32;
 };
 )");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitType_Struct_WithDecoration) {
-  ast::DecorationList decos;
-  decos.push_back(create<ast::StructBlockDecoration>());
-
   auto* s = Structure("S",
                       {
                           Member("a", ty.i32()),
-                          Member("b", ty.f32(), {MemberOffset(4)}),
+                          Member("b", ty.f32(), {MemberAlign(8)}),
                       },
-                      decos);
+                      {create<ast::StructBlockDecoration>()});
 
   GeneratorImpl& gen = Build();
 
@@ -193,7 +250,7 @@ TEST_F(WgslGeneratorImplTest, EmitType_Struct_WithDecoration) {
   EXPECT_EQ(gen.result(), R"([[block]]
 struct S {
   a : i32;
-  [[offset(4)]]
+  [[align(8)]]
   b : f32;
 };
 )");