wsgl parser: pipeline_stage() -> expect_pipeline_stage()

The only place that calls `pipeline_stage()` expects a stage to exist, so follow the `expect_` pattern.

Bug: tint:282
Change-Id: Ie18d24ed25a5f882e66a8e553e53b4fb52dcf6fa
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/31734
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: dan sinclair <dsinclair@chromium.org>

src/reader/wgsl/parser_impl.cc

@@ -1669,14 +1669,10 @@ std::unique_ptr<ast::FunctionDecoration> ParserImpl::function_decoration() {
     if (!expect("stage decoration", Token::Type::kParenLeft))
       return nullptr;
 
-    auto stage = pipeline_stage();
-    if (has_error()) {
+    ast::PipelineStage stage;
+    std::tie(stage, source) = expect_pipeline_stage();
+    if (stage == ast::PipelineStage::kNone)
       return nullptr;
-    }
-    if (stage == ast::PipelineStage::kNone) {
-      add_error(peek(), "invalid value for stage decoration");
-      return nullptr;
-    }
 
     if (!expect("stage decoration", Token::Type::kParenRight))
       return nullptr;
@@ -1784,21 +1780,20 @@ ast::VariableList ParserImpl::param_list() {
 //   : VERTEX
 //   | FRAGMENT
 //   | COMPUTE
-ast::PipelineStage ParserImpl::pipeline_stage() {
+std::pair<ast::PipelineStage, Source> ParserImpl::expect_pipeline_stage() {
+  Source source;
+  if (match(Token::Type::kVertex, &source))
+    return {ast::PipelineStage::kVertex, source};
+
+  if (match(Token::Type::kFragment, &source))
+    return {ast::PipelineStage::kFragment, source};
+
+  if (match(Token::Type::kCompute, &source))
+    return {ast::PipelineStage::kCompute, source};
+
   auto t = peek();
-  if (t.IsVertex()) {
-    next();  // consume the peek
-    return ast::PipelineStage::kVertex;
-  }
-  if (t.IsFragment()) {
-    next();  // consume the peek
-    return ast::PipelineStage::kFragment;
-  }
-  if (t.IsCompute()) {
-    next();  // consume the peek
-    return ast::PipelineStage::kCompute;
-  }
-  return ast::PipelineStage::kNone;
+  add_error(t, "invalid value for stage decoration");
+  return {ast::PipelineStage::kNone, t.source()};
 }
 
 // body_stmt

src/reader/wgsl/parser_impl.h

@@ -257,9 +257,11 @@ class ParserImpl {
   /// Parses a `param_list` grammar element
   /// @returns the parsed variables
   ast::VariableList param_list();
-  /// Parses a `pipeline_stage` grammar element
-  /// @returns the pipeline stage or PipelineStage::kNone if none matched
-  ast::PipelineStage pipeline_stage();
+  /// Parses a `pipeline_stage` grammar element, erroring if the next token does
+  /// not match a stage name.
+  /// @returns the pipeline stage or PipelineStage::kNone if none matched, along
+  /// with the source location for the stage.
+  std::pair<ast::PipelineStage, Source> expect_pipeline_stage();
   /// Parses a `body_stmt` grammar element
   /// @returns the parsed statements
   std::unique_ptr<ast::BlockStatement> body_stmt();

src/reader/wgsl/parser_impl_pipeline_stage_test.cc

@@ -23,12 +23,11 @@ namespace wgsl {
 namespace {
 
 struct PipelineStageData {
-  const char* input;
+  std::string input;
   ast::PipelineStage result;
 };
 inline std::ostream& operator<<(std::ostream& out, PipelineStageData data) {
-  out << std::string(data.input);
-  return out;
+  return out << data.input;
 }
 
 class PipelineStageTest : public ParserImplTestWithParam<PipelineStageData> {};
@@ -37,9 +36,15 @@ TEST_P(PipelineStageTest, Parses) {
   auto params = GetParam();
   auto* p = parser(params.input);
 
-  auto stage = p->pipeline_stage();
-  ASSERT_FALSE(p->has_error());
+  ast::PipelineStage stage;
+  Source source;
+  std::tie(stage, source) = p->expect_pipeline_stage();
+  ASSERT_FALSE(p->has_error()) << p->error();
   EXPECT_EQ(stage, params.result);
+  EXPECT_EQ(source.range.begin.line, 1u);
+  EXPECT_EQ(source.range.begin.column, 1u);
+  EXPECT_EQ(source.range.end.line, 1u);
+  EXPECT_EQ(source.range.end.column, 1u + params.input.size());
 
   auto t = p->next();
   EXPECT_TRUE(t.IsEof());
@@ -54,12 +59,16 @@ INSTANTIATE_TEST_SUITE_P(
 
 TEST_F(ParserImplTest, PipelineStage_NoMatch) {
   auto* p = parser("not-a-stage");
-  auto stage = p->pipeline_stage();
+  ast::PipelineStage stage;
+  Source source;
+  std::tie(stage, source) = p->expect_pipeline_stage();
+  ASSERT_TRUE(p->has_error());
+  ASSERT_EQ(p->error(), "1:1: invalid value for stage decoration");
   ASSERT_EQ(stage, ast::PipelineStage::kNone);
-
-  auto t = p->next();
-  EXPECT_TRUE(t.IsIdentifier());
-  EXPECT_EQ(t.to_str(), "not");
+  EXPECT_EQ(source.range.begin.line, 1u);
+  EXPECT_EQ(source.range.begin.column, 1u);
+  EXPECT_EQ(source.range.end.line, 1u);
+  EXPECT_EQ(source.range.end.column, 4u);
 }
 
 }  // namespace