Add tokens for left and right shift

Manually split `>>` and `>=` tokens when looking for a `>` to
correctly parse ptr/array/vec declarations and initializations.

Bug: tint:171, tint:355
Change-Id: Iee89a844fd999e337ae44ef9b192cc122fbf9e54
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/39362
Auto-Submit: James Price <jrprice@google.com>
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: dan sinclair <dsinclair@chromium.org>

src/reader/wgsl/lexer.cc

@@ -431,6 +431,10 @@ Token Lexer::try_punctuation() {
     type = Token::Type::kGreaterThanEqual;
     pos_ += 2;
     location_.column += 2;
+  } else if (matches(pos_, ">>")) {
+    type = Token::Type::kShiftRight;
+    pos_ += 2;
+    location_.column += 2;
   } else if (matches(pos_, ">")) {
     type = Token::Type::kGreaterThan;
     pos_ += 1;
@@ -439,6 +443,10 @@ Token Lexer::try_punctuation() {
     type = Token::Type::kLessThanEqual;
     pos_ += 2;
     location_.column += 2;
+  } else if (matches(pos_, "<<")) {
+    type = Token::Type::kShiftLeft;
+    pos_ += 2;
+    location_.column += 2;
   } else if (matches(pos_, "<")) {
     type = Token::Type::kLessThan;
     pos_ += 1;

src/reader/wgsl/lexer_test.cc

@@ -430,8 +430,10 @@ INSTANTIATE_TEST_SUITE_P(
                     TokenData{"==", Token::Type::kEqualEqual},
                     TokenData{">", Token::Type::kGreaterThan},
                     TokenData{">=", Token::Type::kGreaterThanEqual},
+                    TokenData{">>", Token::Type::kShiftRight},
                     TokenData{"<", Token::Type::kLessThan},
                     TokenData{"<=", Token::Type::kLessThanEqual},
+                    TokenData{"<<", Token::Type::kShiftLeft},
                     TokenData{"%", Token::Type::kMod},
                     TokenData{"!=", Token::Type::kNotEqual},
                     TokenData{"-", Token::Type::kMinus},

src/reader/wgsl/parser_impl.cc

@@ -2321,23 +2321,20 @@ Maybe<ast::Expression*> ParserImpl::additive_expression() {
 
 // shift_expr
 //   :
-//   | LESS_THAN LESS_THAN additive_expression shift_expr
-//   | GREATER_THAN GREATER_THAN additive_expression shift_expr
+//   | SHIFT_LEFT additive_expression shift_expr
+//   | SHIFT_RIGHT additive_expression shift_expr
 Expect<ast::Expression*> ParserImpl::expect_shift_expr(ast::Expression* lhs) {
   auto t = peek();
   auto source = t.source();
-  auto t2 = peek(1);
 
   auto* name = "";
   ast::BinaryOp op = ast::BinaryOp::kNone;
-  if (t.IsLessThan() && t2.IsLessThan()) {
-    next();  // Consume the t peek
-    next();  // Consume the t2 peek
+  if (t.IsShiftLeft()) {
+    next();  // Consume the peek
     op = ast::BinaryOp::kShiftLeft;
     name = "<<";
-  } else if (t.IsGreaterThan() && t2.IsGreaterThan()) {
-    next();  // Consume the t peek
-    next();  // Consume the t2 peek
+  } else if (t.IsShiftRight()) {
+    next();  // Consume the peek
     op = ast::BinaryOp::kShiftRight;
     name = ">>";
   } else {
@@ -3029,6 +3026,23 @@ bool ParserImpl::expect(const std::string& use, Token::Type tok) {
     return true;
   }
 
+  // Special case to split `>>` and `>=` tokens if we are looking for a `>`.
+  if (tok == Token::Type::kGreaterThan &&
+      (t.IsShiftRight() || t.IsGreaterThanEqual())) {
+    next();
+
+    // Push the second character to the token queue.
+    auto source = t.source();
+    source.range.begin.column++;
+    if (t.IsShiftRight())
+      token_queue_.push_front(Token(Token::Type::kGreaterThan, source));
+    else if (t.IsGreaterThanEqual())
+      token_queue_.push_front(Token(Token::Type::kEqual, source));
+
+    synchronized_ = true;
+    return true;
+  }
+
   std::stringstream err;
   err << "expected '" << Token::TypeToName(tok) << "'";
   if (!use.empty()) {

src/reader/wgsl/parser_impl_shift_expression_test.cc

@@ -71,6 +71,24 @@ TEST_F(ParserImplTest, ShiftExpression_Parses_ShiftRight) {
   ASSERT_TRUE(init->literal()->As<ast::BoolLiteral>()->IsTrue());
 }
 
+TEST_F(ParserImplTest, ShiftExpression_InvalidSpaceLeft) {
+  auto p = parser("a < < true");
+  auto e = p->shift_expression();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  ASSERT_NE(e.value, nullptr);
+  EXPECT_FALSE(e.value->Is<ast::BinaryExpression>());
+}
+
+TEST_F(ParserImplTest, ShiftExpression_InvalidSpaceRight) {
+  auto p = parser("a > > true");
+  auto e = p->shift_expression();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  ASSERT_NE(e.value, nullptr);
+  EXPECT_FALSE(e.value->Is<ast::BinaryExpression>());
+}
+
 TEST_F(ParserImplTest, ShiftExpression_InvalidLHS) {
   auto p = parser("if (a) {} << true");
   auto e = p->shift_expression();

src/reader/wgsl/parser_impl_type_decl_test.cc

@@ -532,6 +532,20 @@ TEST_F(ParserImplTest, TypeDecl_Array_Runtime) {
   ASSERT_TRUE(a->type()->Is<type::U32>());
 }
 
+TEST_F(ParserImplTest, TypeDecl_Array_Runtime_Vec) {
+  auto p = parser("array<vec4<u32>>");
+  auto t = p->type_decl();
+  EXPECT_TRUE(t.matched);
+  EXPECT_FALSE(t.errored);
+  ASSERT_NE(t.value, nullptr) << p->error();
+  ASSERT_FALSE(p->has_error());
+  ASSERT_TRUE(t->Is<type::Array>());
+
+  auto* a = t->As<type::Array>();
+  ASSERT_TRUE(a->IsRuntimeArray());
+  ASSERT_TRUE(a->type()->is_unsigned_integer_vector());
+}
+
 TEST_F(ParserImplTest, TypeDecl_Array_BadType) {
   auto p = parser("array<unknown, 3>");
   auto t = p->type_decl();

src/reader/wgsl/parser_impl_variable_stmt_test.cc

@@ -82,6 +82,66 @@ TEST_F(ParserImplTest, VariableStmt_VariableDecl_ConstructorInvalid) {
   EXPECT_EQ(p->error(), "1:15: missing constructor for variable declaration");
 }
 
+TEST_F(ParserImplTest, VariableStmt_VariableDecl_ArrayInit) {
+  auto p = parser("var a : array<i32> = array<i32>();");
+  auto e = p->variable_stmt();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  EXPECT_FALSE(p->has_error()) << p->error();
+  ASSERT_NE(e.value, nullptr);
+  ASSERT_TRUE(e->Is<ast::VariableDeclStatement>());
+  ASSERT_NE(e->variable(), nullptr);
+  EXPECT_EQ(e->variable()->symbol(), p->builder().Symbols().Get("a"));
+
+  ASSERT_NE(e->variable()->constructor(), nullptr);
+  EXPECT_TRUE(e->variable()->constructor()->Is<ast::ConstructorExpression>());
+}
+
+TEST_F(ParserImplTest, VariableStmt_VariableDecl_ArrayInit_NoSpace) {
+  auto p = parser("var a : array<i32>=array<i32>();");
+  auto e = p->variable_stmt();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  EXPECT_FALSE(p->has_error()) << p->error();
+  ASSERT_NE(e.value, nullptr);
+  ASSERT_TRUE(e->Is<ast::VariableDeclStatement>());
+  ASSERT_NE(e->variable(), nullptr);
+  EXPECT_EQ(e->variable()->symbol(), p->builder().Symbols().Get("a"));
+
+  ASSERT_NE(e->variable()->constructor(), nullptr);
+  EXPECT_TRUE(e->variable()->constructor()->Is<ast::ConstructorExpression>());
+}
+
+TEST_F(ParserImplTest, VariableStmt_VariableDecl_VecInit) {
+  auto p = parser("var a : vec2<i32> = vec2<i32>();");
+  auto e = p->variable_stmt();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  EXPECT_FALSE(p->has_error()) << p->error();
+  ASSERT_NE(e.value, nullptr);
+  ASSERT_TRUE(e->Is<ast::VariableDeclStatement>());
+  ASSERT_NE(e->variable(), nullptr);
+  EXPECT_EQ(e->variable()->symbol(), p->builder().Symbols().Get("a"));
+
+  ASSERT_NE(e->variable()->constructor(), nullptr);
+  EXPECT_TRUE(e->variable()->constructor()->Is<ast::ConstructorExpression>());
+}
+
+TEST_F(ParserImplTest, VariableStmt_VariableDecl_VecInit_NoSpace) {
+  auto p = parser("var a : vec2<i32>=vec2<i32>();");
+  auto e = p->variable_stmt();
+  EXPECT_TRUE(e.matched);
+  EXPECT_FALSE(e.errored);
+  EXPECT_FALSE(p->has_error()) << p->error();
+  ASSERT_NE(e.value, nullptr);
+  ASSERT_TRUE(e->Is<ast::VariableDeclStatement>());
+  ASSERT_NE(e->variable(), nullptr);
+  EXPECT_EQ(e->variable()->symbol(), p->builder().Symbols().Get("a"));
+
+  ASSERT_NE(e->variable()->constructor(), nullptr);
+  EXPECT_TRUE(e->variable()->constructor()->Is<ast::ConstructorExpression>());
+}
+
 TEST_F(ParserImplTest, VariableStmt_Const) {
   auto p = parser("const a : i32 = 1");
   auto e = p->variable_stmt();

src/reader/wgsl/token.cc

@@ -74,10 +74,14 @@ std::string Token::TypeToName(Type type) {
       return ">";
     case Token::Type::kGreaterThanEqual:
       return ">=";
+    case Token::Type::kShiftRight:
+      return ">>";
     case Token::Type::kLessThan:
       return "<";
     case Token::Type::kLessThanEqual:
       return "<=";
+    case Token::Type::kShiftLeft:
+      return "<<";
     case Token::Type::kMod:
       return "%";
     case Token::Type::kNotEqual:

src/reader/wgsl/token.h

@@ -85,10 +85,14 @@ class Token {
     kGreaterThan,
     /// A '>='
     kGreaterThanEqual,
+    /// A '>>'
+    kShiftRight,
     /// A '<'
     kLessThan,
     /// A '<='
     kLessThanEqual,
+    /// A '<<'
+    kShiftLeft,
     /// A '%'
     kMod,
     /// A '-'
@@ -424,10 +428,14 @@ class Token {
   bool IsGreaterThan() const { return type_ == Type::kGreaterThan; }
   /// @returns true if token is a '>='
   bool IsGreaterThanEqual() const { return type_ == Type::kGreaterThanEqual; }
+  /// @returns true if token is a '>>'
+  bool IsShiftRight() const { return type_ == Type::kShiftRight; }
   /// @returns true if token is a '<'
   bool IsLessThan() const { return type_ == Type::kLessThan; }
   /// @returns true if token is a '<='
   bool IsLessThanEqual() const { return type_ == Type::kLessThanEqual; }
+  /// @returns true if token is a '<<'
+  bool IsShiftLeft() const { return type_ == Type::kShiftLeft; }
   /// @returns true if token is a '%'
   bool IsMod() const { return type_ == Type::kMod; }
   /// @returns true if token is a '-'