[wgsl-reader] Add stride support.

This CL adds stride support to the WGSL reader.

Bug: tint:178
Change-Id: Id6b5163438e562a371255ad3fb992d0a716543e7
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/26040
Reviewed-by: David Neto <dneto@google.com>

src/reader/wgsl/lexer.cc

@@ -563,6 +563,8 @@ Token Lexer::check_keyword(const Source& source, const std::string& str) {
     return {Token::Type::kSet, source, "set"};
   if (str == "storage_buffer")
     return {Token::Type::kStorageBuffer, source, "storage_buffer"};
+  if (str == "stride")
+    return {Token::Type::kStride, source, "stride"};
   if (str == "struct")
     return {Token::Type::kStruct, source, "struct"};
   if (str == "switch")

src/reader/wgsl/lexer_test.cc

@@ -458,6 +458,7 @@ INSTANTIATE_TEST_SUITE_P(
                     TokenData{"return", Token::Type::kReturn},
                     TokenData{"set", Token::Type::kSet},
                     TokenData{"storage_buffer", Token::Type::kStorageBuffer},
+                    TokenData{"stride", Token::Type::kStride},
                     TokenData{"struct", Token::Type::kStruct},
                     TokenData{"switch", Token::Type::kSwitch},
                     TokenData{"true", Token::Type::kTrue},

src/reader/wgsl/parser_impl.cc

@@ -686,8 +686,10 @@ ast::type::AliasType* ParserImpl::type_alias() {
 //   | VEC3 LESS_THAN type_decl GREATER_THAN
 //   | VEC4 LESS_THAN type_decl GREATER_THAN
 //   | PTR LESS_THAN storage_class, type_decl GREATER_THAN
-//   | ARRAY LESS_THAN type_decl COMMA INT_LITERAL GREATER_THAN
-//   | ARRAY LESS_THAN type_decl GREATER_THAN
+//   | array_decoration_list? ARRAY LESS_THAN type_decl COMMA
+//          INT_LITERAL GREATER_THAN
+//   | array_decoration_list? ARRAY LESS_THAN type_decl
+//          GREATER_THAN
 //   | MAT2x2 LESS_THAN type_decl GREATER_THAN
 //   | MAT2x3 LESS_THAN type_decl GREATER_THAN
 //   | MAT2x4 LESS_THAN type_decl GREATER_THAN
@@ -730,8 +732,20 @@ ast::type::Type* ParserImpl::type_decl() {
   if (t.IsPtr()) {
     return type_decl_pointer(t);
   }
+
+  auto deco = array_decoration_list();
+  if (has_error()) {
+    return nullptr;
+  }
+  if (deco != 0) {
+    t = peek();
+  }
+  if (deco != 0 && !t.IsArray()) {
+    set_error(t, "found array decoration but no array");
+    return nullptr;
+  }
   if (t.IsArray()) {
-    return type_decl_array(t);
+    return type_decl_array(t, deco);
   }
   if (t.IsMat2x2() || t.IsMat2x3() || t.IsMat2x4() || t.IsMat3x2() ||
       t.IsMat3x3() || t.IsMat3x4() || t.IsMat4x2() || t.IsMat4x3() ||
@@ -815,7 +829,7 @@ ast::type::Type* ParserImpl::type_decl_vector(Token t) {
       std::make_unique<ast::type::VectorType>(subtype, count));
 }
 
-ast::type::Type* ParserImpl::type_decl_array(Token t) {
+ast::type::Type* ParserImpl::type_decl_array(Token t, uint32_t stride) {
   next();  // Consume the peek
 
   t = next();
@@ -852,8 +866,50 @@ ast::type::Type* ParserImpl::type_decl_array(Token t) {
     return nullptr;
   }
 
-  return ctx_.type_mgr().Get(
-      std::make_unique<ast::type::ArrayType>(subtype, size));
+  auto ty = std::make_unique<ast::type::ArrayType>(subtype, size);
+  if (stride != 0) {
+    ty->set_array_stride(stride);
+  }
+  return ctx_.type_mgr().Get(std::move(ty));
+}
+
+// array_decoration_list
+//   : ATTR_LEFT (array_decoration COMMA)* array_decoration ATTR_RIGHT
+// array_decoration
+//   : STRIDE INT_LITERAL
+//
+// As there is currently only one decoration I'm combining these for now.
+// we can split apart later if needed.
+uint32_t ParserImpl::array_decoration_list() {
+  auto t = peek();
+  if (!t.IsAttrLeft()) {
+    return 0;
+  }
+  t = peek(1);
+  if (!t.IsStride()) {
+    return 0;
+  }
+
+  next();  // consume the peek of [[
+  next();  // consume the peek of stride
+
+  t = next();
+  if (!t.IsSintLiteral()) {
+    set_error(t, "missing value for stride decoration");
+    return 0;
+  }
+  if (t.to_i32() < 0) {
+    set_error(t, "invalid stride value: " + t.to_str());
+    return 0;
+  }
+
+  uint32_t stride = static_cast<uint32_t>(t.to_i32());
+  t = next();
+  if (!t.IsAttrRight()) {
+    set_error(t, "missing ]] for array decoration");
+    return 0;
+  }
+  return stride;
 }
 
 ast::type::Type* ParserImpl::type_decl_matrix(Token t) {
@@ -985,16 +1041,16 @@ ast::StructDecoration ParserImpl::struct_decoration_decl() {
   if (!t.IsAttrLeft())
     return ast::StructDecoration::kNone;
 
-  next();  // Consume the peek
-
-  auto deco = struct_decoration();
+  auto deco = struct_decoration(peek(1));
   if (has_error())
     return ast::StructDecoration::kNone;
   if (deco == ast::StructDecoration::kNone) {
-    set_error(peek(), "unknown struct decoration");
-    return ast::StructDecoration::kNone;
+    return deco;
   }
 
+  next();  // Consume the peek of [[
+  next();  // Consume the peek from the struct_decoration
+
   t = next();
   if (!t.IsAttrRight()) {
     set_error(t, "missing ]] for struct decoration");
@@ -1006,10 +1062,8 @@ ast::StructDecoration ParserImpl::struct_decoration_decl() {
 
 // struct_decoration
 //  : BLOCK
-ast::StructDecoration ParserImpl::struct_decoration() {
-  auto t = peek();
+ast::StructDecoration ParserImpl::struct_decoration(Token t) {
   if (t.IsBlock()) {
-    next();  // Consume the peek
     return ast::StructDecoration::kBlock;
   }
   return ast::StructDecoration::kNone;

src/reader/wgsl/parser_impl.h

@@ -144,7 +144,7 @@ class ParserImpl {
   ast::StructDecoration struct_decoration_decl();
   /// Parses a `struct_decoration` grammar element
   /// @returns the struct decoration or StructDecoraton::kNone if none matched
-  ast::StructDecoration struct_decoration();
+  ast::StructDecoration struct_decoration(Token t);
   /// Parses a `struct_body_decl` grammar element
   /// @returns the struct members
   ast::StructMemberList struct_body_decl();
@@ -339,7 +339,8 @@ class ParserImpl {
  private:
   ast::type::Type* type_decl_pointer(Token t);
   ast::type::Type* type_decl_vector(Token t);
-  ast::type::Type* type_decl_array(Token t);
+  ast::type::Type* type_decl_array(Token t, uint32_t stride);
+  uint32_t array_decoration_list();
   ast::type::Type* type_decl_matrix(Token t);
 
   Context& ctx_;

src/reader/wgsl/parser_impl_struct_decoration_decl_test.cc

@@ -35,11 +35,11 @@ TEST_F(ParserImplTest, StructDecorationDecl_MissingAttrRight) {
   EXPECT_EQ(p->error(), "1:8: missing ]] for struct decoration");
 }
 
+// Note, this isn't an error because it could be an array decoration
 TEST_F(ParserImplTest, StructDecorationDecl_InvalidDecoration) {
   auto* p = parser("[[invalid]]");
   p->struct_decoration_decl();
-  ASSERT_TRUE(p->has_error());
-  EXPECT_EQ(p->error(), "1:3: unknown struct decoration");
+  ASSERT_FALSE(p->has_error());
 }
 
 }  // namespace

src/reader/wgsl/parser_impl_struct_decoration_test.cc

@@ -55,12 +55,9 @@ TEST_P(StructDecorationTest, Parses) {
   auto params = GetParam();
   auto* p = parser(params.input);
 
-  auto deco = p->struct_decoration();
+  auto deco = p->struct_decoration(p->peek());
   ASSERT_FALSE(p->has_error());
   EXPECT_EQ(deco, params.result);
-
-  auto t = p->next();
-  EXPECT_TRUE(t.IsEof());
 }
 INSTANTIATE_TEST_SUITE_P(ParserImplTest,
                          StructDecorationTest,
@@ -69,7 +66,7 @@ INSTANTIATE_TEST_SUITE_P(ParserImplTest,
 
 TEST_F(ParserImplTest, StructDecoration_NoMatch) {
   auto* p = parser("not-a-stage");
-  auto deco = p->struct_decoration();
+  auto deco = p->struct_decoration(p->peek());
   ASSERT_EQ(deco, ast::StructDecoration::kNone);
 
   auto t = p->next();

src/reader/wgsl/parser_impl_type_alias_test.cc

@@ -14,6 +14,7 @@
 
 #include "gtest/gtest.h"
 #include "src/ast/type/alias_type.h"
+#include "src/ast/type/array_type.h"
 #include "src/ast/type/i32_type.h"
 #include "src/ast/type/struct_type.h"
 #include "src/reader/wgsl/parser_impl.h"
@@ -88,6 +89,26 @@ TEST_F(ParserImplTest, TypeDecl_InvalidStruct) {
   EXPECT_EQ(p->error(), "1:20: missing struct declaration");
 }
 
+TEST_F(ParserImplTest, TypeDecl_Struct_WithStride) {
+  auto* p = parser(
+      "type a = [[block]] struct { [[offset 0]] data: [[stride 4]] array<f32>; "
+      "}");
+  auto* t = p->type_alias();
+  ASSERT_FALSE(p->has_error());
+  ASSERT_NE(t, nullptr);
+  EXPECT_EQ(t->name(), "a");
+  ASSERT_TRUE(t->type()->IsStruct());
+
+  auto* s = t->type()->AsStruct();
+  EXPECT_EQ(s->impl()->members().size(), 1u);
+
+  const auto* ty = s->impl()->members()[0]->type();
+  ASSERT_TRUE(ty->IsArray());
+  const auto* arr = ty->AsArray();
+  EXPECT_TRUE(arr->has_array_stride());
+  EXPECT_EQ(arr->array_stride(), 4u);
+}
+
 }  // namespace
 }  // namespace wgsl
 }  // namespace reader

src/reader/wgsl/parser_impl_type_decl_test.cc

@@ -389,6 +389,85 @@ TEST_F(ParserImplTest, TypeDecl_Array) {
   ASSERT_FALSE(a->IsRuntimeArray());
   ASSERT_EQ(a->size(), 5u);
   ASSERT_TRUE(a->type()->IsF32());
+  ASSERT_FALSE(a->has_array_stride());
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Stride) {
+  auto* p = parser("[[stride 16]] array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_NE(t, nullptr);
+  ASSERT_FALSE(p->has_error());
+  ASSERT_TRUE(t->IsArray());
+
+  auto* a = t->AsArray();
+  ASSERT_FALSE(a->IsRuntimeArray());
+  ASSERT_EQ(a->size(), 5u);
+  ASSERT_TRUE(a->type()->IsF32());
+  ASSERT_TRUE(a->has_array_stride());
+  EXPECT_EQ(a->array_stride(), 16u);
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Runtime_Stride) {
+  auto* p = parser("[[stride 16]] array<f32>");
+  auto* t = p->type_decl();
+  ASSERT_NE(t, nullptr);
+  ASSERT_FALSE(p->has_error());
+  ASSERT_TRUE(t->IsArray());
+
+  auto* a = t->AsArray();
+  ASSERT_TRUE(a->IsRuntimeArray());
+  ASSERT_TRUE(a->type()->IsF32());
+  ASSERT_TRUE(a->has_array_stride());
+  EXPECT_EQ(a->array_stride(), 16u);
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Decoration_MissingArray) {
+  auto* p = parser("[[stride 16]] f32");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:15: found array decoration but no array");
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Decoration_MissingClosingAttr) {
+  auto* p = parser("[[stride 16 array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:13: missing ]] for array decoration");
+}
+
+// Note, this isn't an error because it could be a struct decoration, we just
+// don't have an array ...
+TEST_F(ParserImplTest, TypeDecl_Array_Decoration_UnknownDecoration) {
+  auto* p = parser("[[unknown 16]] array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_FALSE(p->has_error());
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Stride_MissingValue) {
+  auto* p = parser("[[stride]] array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:9: missing value for stride decoration");
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Stride_InvalidValue) {
+  auto* p = parser("[[stride invalid]] array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:10: missing value for stride decoration");
+}
+
+TEST_F(ParserImplTest, TypeDecl_Array_Stride_InvalidValue_Negative) {
+  auto* p = parser("[[stride -1]] array<f32, 5>");
+  auto* t = p->type_decl();
+  ASSERT_EQ(t, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:10: invalid stride value: -1");
 }
 
 TEST_F(ParserImplTest, TypeDecl_Array_Runtime) {

src/reader/wgsl/token.cc

@@ -199,6 +199,8 @@ std::string Token::TypeToName(Type type) {
       return "set";
     case Token::Type::kStorageBuffer:
       return "storage_buffer";
+    case Token::Type::kStride:
+      return "stride";
     case Token::Type::kStruct:
       return "struct";
     case Token::Type::kSwitch:

src/reader/wgsl/token.h

@@ -210,6 +210,8 @@ class Token {
     kSet,
     /// A 'storage_buffer'
     kStorageBuffer,
+    /// A 'stride'
+    kStride,
     /// A 'struct'
     kStruct,
     /// A 'switch'
@@ -463,6 +465,8 @@ class Token {
   bool IsSet() const { return type_ == Type::kSet; }
   /// @returns true if token is a 'storage_buffer'
   bool IsStorageBuffer() const { return type_ == Type::kStorageBuffer; }
+  /// @returns true if token is a 'stride'
+  bool IsStride() const { return type_ == Type::kStride; }
   /// @returns true if token is a 'struct'
   bool IsStruct() const { return type_ == Type::kStruct; }
   /// @returns true if token is a 'switch'

test/compute_boids.wgsl

@@ -57,7 +57,7 @@ type SimParams = [[block]] struct {
 };
 
 type Particles = [[block]] struct {
-  [[offset 0]] particles : array<Particle, 5>;
+  [[offset 0]] particles : [[stride 16]] array<Particle, 5>;
 };
 
 [[binding 0, set 0]] var<uniform> params : SimParams;