Add support for >>= and <<=.

This CL adds `<<=` and `>>=` to the supported operators in WGSL. The
ExpandCompoundAssignment transform is used to convert to the expanded form.

Bug: tint:1594
Change-Id: I20519052c52d4b69bc90def1acc5c0a30c36fd8a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/97980
Auto-Submit: Dan Sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/tint/reader/wgsl/lexer.cc

@@ -1025,8 +1025,13 @@ Token Lexer::try_punctuation() {
             type = Token::Type::kGreaterThanEqual;
             advance(2);
         } else if (matches(pos() + 1, '>')) {
-            type = Token::Type::kShiftRight;
+            if (matches(pos() + 2, '=')) {
-            advance(2);
+                type = Token::Type::kShiftRightEqual;
+                advance(3);
+            } else {
+                type = Token::Type::kShiftRight;
+                advance(2);
+            }
         } else {
             type = Token::Type::kGreaterThan;
             advance(1);
@@ -1036,8 +1041,13 @@ Token Lexer::try_punctuation() {
             type = Token::Type::kLessThanEqual;
             advance(2);
         } else if (matches(pos() + 1, '<')) {
-            type = Token::Type::kShiftLeft;
+            if (matches(pos() + 2, '=')) {
-            advance(2);
+                type = Token::Type::kShiftLeftEqual;
+                advance(3);
+            } else {
+                type = Token::Type::kShiftLeft;
+                advance(2);
+            }
         } else {
             type = Token::Type::kLessThan;
             advance(1);

src/tint/reader/wgsl/lexer_test.cc

@@ -1049,7 +1049,9 @@ INSTANTIATE_TEST_SUITE_P(LexerTest,
                                          TokenData{"%=", Token::Type::kModuloEqual},
                                          TokenData{"&=", Token::Type::kAndEqual},
                                          TokenData{"|=", Token::Type::kOrEqual},
-                                         TokenData{"^=", Token::Type::kXorEqual}));
+                                         TokenData{"^=", Token::Type::kXorEqual},
+                                         TokenData{">>=", Token::Type::kShiftRightEqual},
+                                         TokenData{"<<=", Token::Type::kShiftLeftEqual}));
 
 using SplittablePunctuationTest = testing::TestWithParam<TokenData>;
 TEST_P(SplittablePunctuationTest, Parses) {

src/tint/reader/wgsl/parser_impl.cc

@@ -2924,6 +2924,8 @@ Maybe<const ast::Expression*> ParserImpl::logical_or_expression() {
 // | and_equal
 // | or_equal
 // | xor_equal
+// | shift_right_equal
+// | shift_left_equal
 Maybe<ast::BinaryOp> ParserImpl::compound_assignment_operator() {
     ast::BinaryOp compound_op = ast::BinaryOp::kNone;
     if (peek_is(Token::Type::kPlusEqual)) {
@@ -2942,6 +2944,10 @@ Maybe<ast::BinaryOp> ParserImpl::compound_assignment_operator() {
         compound_op = ast::BinaryOp::kOr;
     } else if (peek_is(Token::Type::kXorEqual)) {
         compound_op = ast::BinaryOp::kXor;
+    } else if (peek_is(Token::Type::kShiftLeftEqual)) {
+        compound_op = ast::BinaryOp::kShiftLeft;
+    } else if (peek_is(Token::Type::kShiftRightEqual)) {
+        compound_op = ast::BinaryOp::kShiftRight;
     }
     if (compound_op != ast::BinaryOp::kNone) {
         next();

src/tint/reader/wgsl/parser_impl_assignment_stmt_test.cc

@@ -111,8 +111,14 @@ TEST_F(ParserImplTest, AssignmentStmt_Parses_ToPhony) {
     ASSERT_TRUE(a->lhs->Is<ast::PhonyExpression>());
 }
 
-TEST_F(ParserImplTest, AssignmentStmt_Parses_CompoundOp) {
+struct CompoundData {
-    auto p = parser("a += 123u");
+    std::string str;
+    ast::BinaryOp op;
+};
+using CompoundOpTest = ParserImplTestWithParam<CompoundData>;
+TEST_P(CompoundOpTest, CompoundOp) {
+    auto params = GetParam();
+    auto p = parser("a " + params.str + " 123u");
     auto e = p->assignment_stmt();
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
@@ -123,7 +129,7 @@ TEST_F(ParserImplTest, AssignmentStmt_Parses_CompoundOp) {
     ASSERT_NE(a, nullptr);
     ASSERT_NE(a->lhs, nullptr);
     ASSERT_NE(a->rhs, nullptr);
-    EXPECT_EQ(a->op, ast::BinaryOp::kAdd);
+    EXPECT_EQ(a->op, params.op);
 
     ASSERT_TRUE(a->lhs->Is<ast::IdentifierExpression>());
     auto* ident = a->lhs->As<ast::IdentifierExpression>();
@@ -134,6 +140,18 @@ TEST_F(ParserImplTest, AssignmentStmt_Parses_CompoundOp) {
     EXPECT_EQ(a->rhs->As<ast::IntLiteralExpression>()->suffix,
               ast::IntLiteralExpression::Suffix::kU);
 }
+INSTANTIATE_TEST_SUITE_P(ParserImplTest,
+                         CompoundOpTest,
+                         testing::Values(CompoundData{"+=", ast::BinaryOp::kAdd},
+                                         CompoundData{"-=", ast::BinaryOp::kSubtract},
+                                         CompoundData{"*=", ast::BinaryOp::kMultiply},
+                                         CompoundData{"/=", ast::BinaryOp::kDivide},
+                                         CompoundData{"%=", ast::BinaryOp::kModulo},
+                                         CompoundData{"&=", ast::BinaryOp::kAnd},
+                                         CompoundData{"|=", ast::BinaryOp::kOr},
+                                         CompoundData{"^=", ast::BinaryOp::kXor},
+                                         CompoundData{">>=", ast::BinaryOp::kShiftRight},
+                                         CompoundData{"<<=", ast::BinaryOp::kShiftLeft}));
 
 TEST_F(ParserImplTest, AssignmentStmt_MissingEqual) {
     auto p = parser("a.b.c[2].d 123");

src/tint/reader/wgsl/token.cc

@@ -130,6 +130,10 @@ std::string_view Token::TypeToName(Type type) {
             return "|=";
         case Token::Type::kXorEqual:
             return "^=";
+        case Token::Type::kShiftLeftEqual:
+            return "<<=";
+        case Token::Type::kShiftRightEqual:
+            return ">>=";
 
         case Token::Type::kArray:
             return "array";

src/tint/reader/wgsl/token.h

@@ -140,6 +140,10 @@ class Token {
         kOrEqual,
         /// A '^='
         kXorEqual,
+        /// A '>>='
+        kShiftRightEqual,
+        /// A '<<='
+        kShiftLeftEqual,
 
         /// A 'array'
         kArray,

test/tint/statements/compound_assign/scalar/shift_left.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : i32,
+}
+@group(0) @binding(0)
+var<storage, read_write> v : S;
+
+fn foo() {
+  v.a <<= 2;
+}

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.dxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store(0u, asuint((asint(v.Load(0u)) << 2u)));
+}

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.fxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store(0u, asuint((asint(v.Load(0u)) << 2u)));
+}

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.glsl

@@ -0,0 +1,17 @@
+#version 310 es
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+void unused_entry_point() {
+  return;
+}
+struct S {
+  int a;
+};
+
+layout(binding = 0, std430) buffer S_1 {
+  int a;
+} v;
+void foo() {
+  v.a = (v.a << 2u);
+}
+

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.msl

@@ -0,0 +1,11 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct S {
+  /* 0x0000 */ int a;
+};
+
+void foo(device S* const tint_symbol) {
+  (*(tint_symbol)).a = as_type<int>((as_type<uint>((*(tint_symbol)).a) << 2u));
+}
+

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.spvasm

@@ -0,0 +1,41 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 19
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %S "S"
+               OpMemberName %S 0 "a"
+               OpName %v "v"
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %foo "foo"
+               OpDecorate %S Block
+               OpMemberDecorate %S 0 Offset 0
+               OpDecorate %v DescriptorSet 0
+               OpDecorate %v Binding 0
+        %int = OpTypeInt 32 1
+          %S = OpTypeStruct %int
+%_ptr_StorageBuffer_S = OpTypePointer StorageBuffer %S
+          %v = OpVariable %_ptr_StorageBuffer_S StorageBuffer
+       %void = OpTypeVoid
+          %5 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_int = OpTypePointer StorageBuffer %int
+     %uint_2 = OpConstant %uint 2
+%unused_entry_point = OpFunction %void None %5
+          %8 = OpLabel
+               OpReturn
+               OpFunctionEnd
+        %foo = OpFunction %void None %5
+         %10 = OpLabel
+         %14 = OpAccessChain %_ptr_StorageBuffer_int %v %uint_0
+         %15 = OpAccessChain %_ptr_StorageBuffer_int %v %uint_0
+         %16 = OpLoad %int %15
+         %18 = OpShiftLeftLogical %int %16 %uint_2
+               OpStore %14 %18
+               OpReturn
+               OpFunctionEnd

test/tint/statements/compound_assign/scalar/shift_left.wgsl.expected.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : i32,
+}
+
+@group(0) @binding(0) var<storage, read_write> v : S;
+
+fn foo() {
+  v.a <<= 2;
+}

test/tint/statements/compound_assign/scalar/shift_right.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : i32,
+}
+@group(0) @binding(0)
+var<storage, read_write> v : S;
+
+fn foo() {
+  v.a >>= 2;
+}

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.dxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store(0u, asuint((asint(v.Load(0u)) >> 2u)));
+}

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.fxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store(0u, asuint((asint(v.Load(0u)) >> 2u)));
+}

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.glsl

@@ -0,0 +1,17 @@
+#version 310 es
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+void unused_entry_point() {
+  return;
+}
+struct S {
+  int a;
+};
+
+layout(binding = 0, std430) buffer S_1 {
+  int a;
+} v;
+void foo() {
+  v.a = (v.a >> 2u);
+}
+

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.msl

@@ -0,0 +1,11 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct S {
+  /* 0x0000 */ int a;
+};
+
+void foo(device S* const tint_symbol) {
+  (*(tint_symbol)).a = ((*(tint_symbol)).a >> 2u);
+}
+

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.spvasm

@@ -0,0 +1,41 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 19
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %S "S"
+               OpMemberName %S 0 "a"
+               OpName %v "v"
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %foo "foo"
+               OpDecorate %S Block
+               OpMemberDecorate %S 0 Offset 0
+               OpDecorate %v DescriptorSet 0
+               OpDecorate %v Binding 0
+        %int = OpTypeInt 32 1
+          %S = OpTypeStruct %int
+%_ptr_StorageBuffer_S = OpTypePointer StorageBuffer %S
+          %v = OpVariable %_ptr_StorageBuffer_S StorageBuffer
+       %void = OpTypeVoid
+          %5 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_int = OpTypePointer StorageBuffer %int
+     %uint_2 = OpConstant %uint 2
+%unused_entry_point = OpFunction %void None %5
+          %8 = OpLabel
+               OpReturn
+               OpFunctionEnd
+        %foo = OpFunction %void None %5
+         %10 = OpLabel
+         %14 = OpAccessChain %_ptr_StorageBuffer_int %v %uint_0
+         %15 = OpAccessChain %_ptr_StorageBuffer_int %v %uint_0
+         %16 = OpLoad %int %15
+         %18 = OpShiftRightArithmetic %int %16 %uint_2
+               OpStore %14 %18
+               OpReturn
+               OpFunctionEnd

test/tint/statements/compound_assign/scalar/shift_right.wgsl.expected.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : i32,
+}
+
+@group(0) @binding(0) var<storage, read_write> v : S;
+
+fn foo() {
+  v.a >>= 2;
+}

test/tint/statements/compound_assign/vector/shift_left.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : vec4<i32>,
+}
+@group(0) @binding(0)
+var<storage, read_write> v : S;
+
+fn foo() {
+  v.a <<= vec4<u32>(2);
+}

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.dxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store4(0u, asuint((asint(v.Load4(0u)) << (2u).xxxx)));
+}

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.fxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store4(0u, asuint((asint(v.Load4(0u)) << (2u).xxxx)));
+}

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.glsl

@@ -0,0 +1,17 @@
+#version 310 es
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+void unused_entry_point() {
+  return;
+}
+struct S {
+  ivec4 a;
+};
+
+layout(binding = 0, std430) buffer S_1 {
+  ivec4 a;
+} v;
+void foo() {
+  v.a = (v.a << uvec4(2u));
+}
+

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.msl

@@ -0,0 +1,11 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct S {
+  /* 0x0000 */ int4 a;
+};
+
+void foo(device S* const tint_symbol) {
+  (*(tint_symbol)).a = as_type<int4>((as_type<uint4>((*(tint_symbol)).a) << uint4(2u)));
+}
+

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.spvasm

@@ -0,0 +1,44 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 22
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %S "S"
+               OpMemberName %S 0 "a"
+               OpName %v "v"
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %foo "foo"
+               OpDecorate %S Block
+               OpMemberDecorate %S 0 Offset 0
+               OpDecorate %v DescriptorSet 0
+               OpDecorate %v Binding 0
+        %int = OpTypeInt 32 1
+      %v4int = OpTypeVector %int 4
+          %S = OpTypeStruct %v4int
+%_ptr_StorageBuffer_S = OpTypePointer StorageBuffer %S
+          %v = OpVariable %_ptr_StorageBuffer_S StorageBuffer
+       %void = OpTypeVoid
+          %6 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_v4int = OpTypePointer StorageBuffer %v4int
+     %v4uint = OpTypeVector %uint 4
+     %uint_2 = OpConstant %uint 2
+         %20 = OpConstantComposite %v4uint %uint_2 %uint_2 %uint_2 %uint_2
+%unused_entry_point = OpFunction %void None %6
+          %9 = OpLabel
+               OpReturn
+               OpFunctionEnd
+        %foo = OpFunction %void None %6
+         %11 = OpLabel
+         %15 = OpAccessChain %_ptr_StorageBuffer_v4int %v %uint_0
+         %16 = OpAccessChain %_ptr_StorageBuffer_v4int %v %uint_0
+         %17 = OpLoad %v4int %16
+         %21 = OpShiftLeftLogical %v4int %17 %20
+               OpStore %15 %21
+               OpReturn
+               OpFunctionEnd

test/tint/statements/compound_assign/vector/shift_left.wgsl.expected.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : vec4<i32>,
+}
+
+@group(0) @binding(0) var<storage, read_write> v : S;
+
+fn foo() {
+  v.a <<= vec4<u32>(2);
+}

test/tint/statements/compound_assign/vector/shift_right.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : vec4<i32>,
+}
+@group(0) @binding(0)
+var<storage, read_write> v : S;
+
+fn foo() {
+  v.a >>= vec4<u32>(2);
+}

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.dxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store4(0u, asuint((asint(v.Load4(0u)) >> (2u).xxxx)));
+}

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.fxc.hlsl

@@ -0,0 +1,10 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+RWByteAddressBuffer v : register(u0, space0);
+
+void foo() {
+  v.Store4(0u, asuint((asint(v.Load4(0u)) >> (2u).xxxx)));
+}

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.glsl

@@ -0,0 +1,17 @@
+#version 310 es
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+void unused_entry_point() {
+  return;
+}
+struct S {
+  ivec4 a;
+};
+
+layout(binding = 0, std430) buffer S_1 {
+  ivec4 a;
+} v;
+void foo() {
+  v.a = (v.a >> uvec4(2u));
+}
+

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.msl

@@ -0,0 +1,11 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct S {
+  /* 0x0000 */ int4 a;
+};
+
+void foo(device S* const tint_symbol) {
+  (*(tint_symbol)).a = ((*(tint_symbol)).a >> uint4(2u));
+}
+

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.spvasm

@@ -0,0 +1,44 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 22
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %S "S"
+               OpMemberName %S 0 "a"
+               OpName %v "v"
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %foo "foo"
+               OpDecorate %S Block
+               OpMemberDecorate %S 0 Offset 0
+               OpDecorate %v DescriptorSet 0
+               OpDecorate %v Binding 0
+        %int = OpTypeInt 32 1
+      %v4int = OpTypeVector %int 4
+          %S = OpTypeStruct %v4int
+%_ptr_StorageBuffer_S = OpTypePointer StorageBuffer %S
+          %v = OpVariable %_ptr_StorageBuffer_S StorageBuffer
+       %void = OpTypeVoid
+          %6 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_v4int = OpTypePointer StorageBuffer %v4int
+     %v4uint = OpTypeVector %uint 4
+     %uint_2 = OpConstant %uint 2
+         %20 = OpConstantComposite %v4uint %uint_2 %uint_2 %uint_2 %uint_2
+%unused_entry_point = OpFunction %void None %6
+          %9 = OpLabel
+               OpReturn
+               OpFunctionEnd
+        %foo = OpFunction %void None %6
+         %11 = OpLabel
+         %15 = OpAccessChain %_ptr_StorageBuffer_v4int %v %uint_0
+         %16 = OpAccessChain %_ptr_StorageBuffer_v4int %v %uint_0
+         %17 = OpLoad %v4int %16
+         %21 = OpShiftRightArithmetic %v4int %17 %20
+               OpStore %15 %21
+               OpReturn
+               OpFunctionEnd

test/tint/statements/compound_assign/vector/shift_right.wgsl.expected.wgsl

@@ -0,0 +1,9 @@
+struct S {
+  a : vec4<i32>,
+}
+
+@group(0) @binding(0) var<storage, read_write> v : S;
+
+fn foo() {
+  v.a >>= vec4<u32>(2);
+}