[wgsl-reader] Parse stage decoration

This CL adds parsing of the stage decoration to the WGSL parser. The
decoration is not hooked up yet so it's effectively ignored.

Change-Id: I8d86c55cee189f993c10b6da31a9c388ba452021
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/28664
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: David Neto <dneto@google.com>

BUILD.gn

@@ -301,6 +301,8 @@ source_set("libtint_core_src") {
     "src/ast/set_decoration.h",
     "src/ast/sint_literal.cc",
     "src/ast/sint_literal.h",
+    "src/ast/stage_decoration.cc",
+    "src/ast/stage_decoration.h",
     "src/ast/statement.cc",
     "src/ast/statement.h",
     "src/ast/storage_class.cc",
@@ -721,6 +723,7 @@ source_set("tint_unittests_core_src") {
     "src/ast/scalar_constructor_expression_test.cc",
     "src/ast/set_decoration_test.cc",
     "src/ast/sint_literal_test.cc",
+    "src/ast/stage_decoration_test.cc",
     "src/ast/struct_member_offset_decoration_test.cc",
     "src/ast/struct_member_test.cc",
     "src/ast/struct_test.cc",

src/CMakeLists.txt

@@ -122,6 +122,8 @@ set(TINT_LIB_SRCS
   ast/set_decoration.h
   ast/sint_literal.cc
   ast/sint_literal.h
+  ast/stage_decoration.cc
+  ast/stage_decoration.h
   ast/statement.cc
   ast/statement.h
   ast/storage_class.cc
@@ -331,6 +333,7 @@ set(TINT_TEST_SRCS
   ast/scalar_constructor_expression_test.cc
   ast/set_decoration_test.cc
   ast/sint_literal_test.cc
+  ast/stage_decoration_test.cc
   ast/struct_member_test.cc
   ast/struct_member_offset_decoration_test.cc
   ast/struct_test.cc

src/ast/function_decoration.cc

@@ -16,6 +16,7 @@
 
 #include <assert.h>
 
+#include "src/ast/stage_decoration.h"
 #include "src/ast/workgroup_decoration.h"
 
 namespace tint {
@@ -25,10 +26,19 @@ FunctionDecoration::FunctionDecoration() = default;
 
 FunctionDecoration::~FunctionDecoration() = default;
 
+bool FunctionDecoration::IsStage() const {
+  return false;
+}
+
 bool FunctionDecoration::IsWorkgroup() const {
   return false;
 }
 
+const StageDecoration* FunctionDecoration::AsStage() const {
+  assert(IsStage());
+  return static_cast<const StageDecoration*>(this);
+}
+
 const WorkgroupDecoration* FunctionDecoration::AsWorkgroup() const {
   assert(IsWorkgroup());
   return static_cast<const WorkgroupDecoration*>(this);

src/ast/function_decoration.h

@@ -22,6 +22,7 @@
 namespace tint {
 namespace ast {
 
+class StageDecoration;
 class WorkgroupDecoration;
 
 /// A decoration attached to a function
@@ -29,9 +30,13 @@ class FunctionDecoration {
  public:
   virtual ~FunctionDecoration();
 
+  /// @returns true if this is a stage decoration
+  virtual bool IsStage() const;
   /// @returns true if this is a workgroup decoration
   virtual bool IsWorkgroup() const;
 
+  /// @returns the decoration as a stage decoration
+  const StageDecoration* AsStage() const;
   /// @returns the decoration as a workgroup decoration
   const WorkgroupDecoration* AsWorkgroup() const;
 

src/ast/stage_decoration.cc

@@ -0,0 +1,33 @@
+// Copyright 2020 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/ast/stage_decoration.h"
+
+namespace tint {
+namespace ast {
+
+StageDecoration::StageDecoration(ast::PipelineStage stage) : stage_(stage) {}
+
+StageDecoration::~StageDecoration() = default;
+
+bool StageDecoration::IsStage() const {
+  return true;
+}
+
+void StageDecoration::to_str(std::ostream& out) const {
+  out << "StageDecoration{" << stage_ << "}" << std::endl;
+}
+
+}  // namespace ast
+}  // namespace tint

src/ast/stage_decoration.h

@@ -0,0 +1,49 @@
+// Copyright 2020 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_AST_STAGE_DECORATION_H_
+#define SRC_AST_STAGE_DECORATION_H_
+
+#include "src/ast/function_decoration.h"
+#include "src/ast/pipeline_stage.h"
+
+namespace tint {
+namespace ast {
+
+/// A workgroup decoration
+class StageDecoration : public FunctionDecoration {
+ public:
+  /// constructor
+  /// @param stage the pipeline stage
+  explicit StageDecoration(ast::PipelineStage stage);
+  ~StageDecoration() override;
+
+  /// @returns true if this is a stage decoration
+  bool IsStage() const override;
+
+  /// @returns the stage
+  ast::PipelineStage value() const { return stage_; }
+
+  /// Outputs the decoration to the given stream
+  /// @param out the stream to output too
+  void to_str(std::ostream& out) const override;
+
+ private:
+  ast::PipelineStage stage_ = ast::PipelineStage::kNone;
+};
+
+}  // namespace ast
+}  // namespace tint
+
+#endif  // SRC_AST_STAGE_DECORATION_H_

src/ast/stage_decoration_test.cc

@@ -0,0 +1,48 @@
+// Copyright 2020 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/ast/stage_decoration.h"
+
+#include <sstream>
+
+#include "gtest/gtest.h"
+
+namespace tint {
+namespace ast {
+namespace {
+
+using StageDecorationTest = testing::Test;
+
+TEST_F(StageDecorationTest, Creation_1param) {
+  StageDecoration d{ast::PipelineStage::kFragment};
+  EXPECT_EQ(d.value(), ast::PipelineStage::kFragment);
+}
+
+TEST_F(StageDecorationTest, Is) {
+  StageDecoration d{ast::PipelineStage::kFragment};
+  EXPECT_FALSE(d.IsWorkgroup());
+  EXPECT_TRUE(d.IsStage());
+}
+
+TEST_F(StageDecorationTest, ToStr) {
+  StageDecoration d{ast::PipelineStage::kFragment};
+  std::ostringstream out;
+  d.to_str(out);
+  EXPECT_EQ(out.str(), R"(StageDecoration{fragment}
+)");
+}
+
+}  // namespace
+}  // namespace ast
+}  // namespace tint

src/ast/workgroup_decoration_test.cc

@@ -59,6 +59,7 @@ TEST_F(WorkgroupDecorationTest, Creation_3param) {
 TEST_F(WorkgroupDecorationTest, Is) {
   WorkgroupDecoration d{2, 4, 6};
   EXPECT_TRUE(d.IsWorkgroup());
+  EXPECT_FALSE(d.IsStage());
 }
 
 TEST_F(WorkgroupDecorationTest, ToStr) {

src/reader/wgsl/lexer.cc

@@ -637,6 +637,8 @@ Token Lexer::check_keyword(const Source& source, const std::string& str) {
     return {Token::Type::kComparisonSampler, source, "sampler_comparison"};
   if (str == "set")
     return {Token::Type::kSet, source, "set"};
+  if (str == "stage")
+    return {Token::Type::kStage, source, "stage"};
   if (str == "storage_buffer")
     return {Token::Type::kStorageBuffer, source, "storage_buffer"};
   if (str == "stride")

src/reader/wgsl/lexer_test.cc

@@ -496,6 +496,7 @@ INSTANTIATE_TEST_SUITE_P(
         TokenData{"sampler", Token::Type::kSampler},
         TokenData{"sampler_comparison", Token::Type::kComparisonSampler},
         TokenData{"set", Token::Type::kSet},
+        TokenData{"stage", Token::Type::kStage},
         TokenData{"storage_buffer", Token::Type::kStorageBuffer},
         TokenData{"stride", Token::Type::kStride},
         TokenData{"struct", Token::Type::kStruct},

src/reader/wgsl/parser_impl.cc

@@ -41,6 +41,7 @@
 #include "src/ast/scalar_constructor_expression.h"
 #include "src/ast/set_decoration.h"
 #include "src/ast/sint_literal.h"
+#include "src/ast/stage_decoration.h"
 #include "src/ast/struct_member_offset_decoration.h"
 #include "src/ast/switch_statement.h"
 #include "src/ast/type/alias_type.h"
@@ -113,7 +114,7 @@ bool IsVariableDecoration(Token t) {
 }
 
 bool IsFunctionDecoration(Token t) {
-  return t.IsWorkgroupSize();
+  return t.IsWorkgroupSize() || t.IsStage();
 }
 
 }  // namespace
@@ -1837,7 +1838,7 @@ bool ParserImpl::function_decoration_decl(ast::FunctionDecorationList& decos) {
 }
 
 // function_decoration
-//   : TODO(dsinclair) STAGE PAREN_LEFT pipeline_stage PAREN_RIGHT
+//   : STAGE PAREN_LEFT pipeline_stage PAREN_RIGHT
 //   | WORKGROUP_SIZE PAREN_LEFT INT_LITERAL
 //         (COMMA INT_LITERAL (COMMA INT_LITERAL)?)? PAREN_RIGHT
 std::unique_ptr<ast::FunctionDecoration> ParserImpl::function_decoration() {
@@ -1905,6 +1906,31 @@ std::unique_ptr<ast::FunctionDecoration> ParserImpl::function_decoration() {
     return std::make_unique<ast::WorkgroupDecoration>(uint32_t(x), uint32_t(y),
                                                       uint32_t(z));
   }
+  if (t.IsStage()) {
+    next();  // Consume the peek
+
+    t = next();
+    if (!t.IsParenLeft()) {
+      set_error(t, "missing ( for stage decoration");
+      return nullptr;
+    }
+
+    auto stage = pipeline_stage();
+    if (has_error()) {
+      return nullptr;
+    }
+    if (stage == ast::PipelineStage::kNone) {
+      set_error(peek(), "invalid value for stage decoration");
+      return nullptr;
+    }
+
+    t = next();
+    if (!t.IsParenRight()) {
+      set_error(t, "missing ) for stage decoration");
+      return nullptr;
+    }
+    return std::make_unique<ast::StageDecoration>(stage);
+  }
   return nullptr;
 }
 

src/reader/wgsl/parser_impl_function_decoration_test.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "gtest/gtest.h"
+#include "src/ast/stage_decoration.h"
 #include "src/ast/workgroup_decoration.h"
 #include "src/reader/wgsl/parser_impl.h"
 #include "src/reader/wgsl/parser_impl_test_helper.h"
@@ -190,6 +191,47 @@ TEST_F(ParserImplTest, FunctionDecoration_Workgroup_Missing_Z_Invalid) {
   EXPECT_EQ(p->error(), "1:22: missing z value for workgroup_size");
 }
 
+TEST_F(ParserImplTest, FunctionDecoration_Stage) {
+  auto* p = parser("stage(compute)");
+  auto deco = p->function_decoration();
+  ASSERT_NE(deco, nullptr);
+  ASSERT_FALSE(p->has_error());
+  ASSERT_TRUE(deco->IsStage());
+  EXPECT_EQ(deco->AsStage()->value(), ast::PipelineStage::kCompute);
+}
+
+TEST_F(ParserImplTest, FunctionDecoration_Stage_MissingValue) {
+  auto* p = parser("stage()");
+  auto deco = p->function_decoration();
+  ASSERT_EQ(deco, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:7: invalid value for stage decoration");
+}
+
+TEST_F(ParserImplTest, FunctionDecoration_Stage_MissingInvalid) {
+  auto* p = parser("stage(nan)");
+  auto deco = p->function_decoration();
+  ASSERT_EQ(deco, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:7: invalid value for stage decoration");
+}
+
+TEST_F(ParserImplTest, FunctionDecoration_Stage_MissingLeftParen) {
+  auto* p = parser("stage compute)");
+  auto deco = p->function_decoration();
+  ASSERT_EQ(deco, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:7: missing ( for stage decoration");
+}
+
+TEST_F(ParserImplTest, FunctionDecoration_Stage_MissingRightParen) {
+  auto* p = parser("stage(compute");
+  auto deco = p->function_decoration();
+  ASSERT_EQ(deco, nullptr);
+  ASSERT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:14: missing ) for stage decoration");
+}
+
 }  // namespace
 }  // namespace wgsl
 }  // namespace reader

src/reader/wgsl/token.cc

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

src/reader/wgsl/token.h

@@ -286,6 +286,8 @@ class Token {
     kSet,
     /// A 'storage_buffer'
     kStorageBuffer,
+    /// A 'stage'
+    kStage,
     /// A 'stride'
     kStride,
     /// A 'struct'
@@ -511,6 +513,8 @@ class Token {
   bool IsCase() const { return type_ == Type::kCase; }
   /// @returns true if token is a 'cast'
   bool IsCast() const { return type_ == Type::kCast; }
+  /// @returns true if token is a 'sampler_comparison'
+  bool IsComparisonSampler() const { return type_ == Type::kComparisonSampler; }
   /// @returns true if token is a 'compute'
   bool IsCompute() const { return type_ == Type::kCompute; }
   /// @returns true if token is a 'const'
@@ -665,10 +669,10 @@ class Token {
   bool IsReturn() const { return type_ == Type::kReturn; }
   /// @returns true if token is a 'sampler'
   bool IsSampler() const { return type_ == Type::kSampler; }
-  /// @returns true if token is a 'sampler_comparison'
-  bool IsComparisonSampler() const { return type_ == Type::kComparisonSampler; }
   /// @returns true if token is a 'set'
   bool IsSet() const { return type_ == Type::kSet; }
+  /// @returns true if token is a 'stage'
+  bool IsStage() const { return type_ == Type::kStage; }
   /// @returns true if token is a 'storage_buffer'
   bool IsStorageBuffer() const { return type_ == Type::kStorageBuffer; }
   /// @returns true if token is a 'stride'

test/function.wgsl

@@ -16,6 +16,7 @@ fn main() -> f32 {
     return ((2. * 3.) - 4.) / 5.;
 }
 
+[[stage(compute)]]
 [[workgroup_size(2)]]
 fn ep() -> void {
   return;