Handle peeking past placeholder elements.

Currently when peeking, if there is placeholder then anything after
the `peek(0)` will be off by 1 position as the placeholder will
end up included in the index count.

This CL updates the peek routine to correctly skip over any
placeholder element between the current index and the requested
peek token.

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

src/tint/reader/wgsl/parser_impl.cc

@@ -271,23 +271,18 @@ const Token& ParserImpl::next() {
     return tokens_[last_source_idx_];
 }
 
-const Token& ParserImpl::peek(size_t idx) {
-    if (next_token_idx_ + idx >= tokens_.size()) {
-        return tokens_[tokens_.size() - 1];
-    }
-
-    // Skip over any placeholder elements
-    while (true) {
-        if (!tokens_[next_token_idx_ + idx].IsPlaceholder()) {
-            break;
+const Token& ParserImpl::peek(size_t count) {
+    for (size_t idx = next_token_idx_; idx < tokens_.size(); idx++) {
+        if (tokens_[idx].IsPlaceholder()) {
+            continue;
         }
-        idx++;
+        if (count == 0) {
+            return tokens_[idx];
+        }
+        count--;
     }
-    if (next_token_idx_ + idx >= tokens_.size()) {
-        return tokens_[tokens_.size() - 1];
-    }
-
-    return tokens_[next_token_idx_ + idx];
+    // Walked off the end of the token list, return last token.
+    return tokens_[tokens_.size() - 1];
 }
 
 bool ParserImpl::peek_is(Token::Type tok, size_t idx) {

src/tint/reader/wgsl/parser_impl.h

@@ -782,6 +782,11 @@ class ParserImpl {
     /// @return the parsed attribute, or nullptr on error.
     Expect<const ast::Attribute*> expect_attribute();
 
+    /// Splits a peekable token into to parts filling in the peekable fields.
+    /// @param lhs the token to set in the current position
+    /// @param rhs the token to set in the placeholder
+    void split_token(Token::Type lhs, Token::Type rhs);
+
   private:
     /// ReturnType resolves to the return type for the function or lambda F.
     template <typename F>
@@ -956,8 +961,6 @@ class ParserImpl {
     Maybe<const ast::Statement*> for_header_initializer();
     Maybe<const ast::Statement*> for_header_continuing();
 
-    void split_token(Token::Type lhs, Token::Type rhs);
-
     class MultiTokenSource;
     MultiTokenSource make_source_range();
     MultiTokenSource make_source_range_from(const Source& start);

src/tint/reader/wgsl/parser_impl_test.cc

@@ -136,5 +136,66 @@ fn main() -> @location(0) vec4<f32> {
     EXPECT_EQ(p->error(), "5:3: unterminated block comment") << p->error();
 }
 
+TEST_F(ParserImplTest, Peek) {
+    auto p = parser("a == if");
+    EXPECT_TRUE(p->peek_is(Token::Type::kIdentifier));
+    EXPECT_TRUE(p->peek_is(Token::Type::kEqualEqual, 1));
+    EXPECT_TRUE(p->peek_is(Token::Type::kIf, 2));
+}
+
+TEST_F(ParserImplTest, Peek_Placeholder) {
+    auto p = parser(">> if");
+    EXPECT_TRUE(p->peek_is(Token::Type::kShiftRight));
+    EXPECT_TRUE(p->peek_is(Token::Type::kIf, 1));
+}
+
+TEST_F(ParserImplTest, Peek_PastPlaceholder) {
+    auto p = parser(">= vec2<u32>");
+    auto& n = p->next();
+    ASSERT_TRUE(n.Is(Token::Type::kGreaterThanEqual));
+    EXPECT_TRUE(p->peek_is(Token::Type::kVec2)) << "expected: vec2 got: " << p->peek().to_name();
+    EXPECT_TRUE(p->peek_is(Token::Type::kLessThan, 1))
+        << "expected: < got: " << p->peek(1).to_name();
+}
+
+TEST_F(ParserImplTest, Peek_MultiplePlaceholder) {
+    auto p = parser(">= >= vec2<u32>");
+    auto& n = p->next();
+    ASSERT_TRUE(n.Is(Token::Type::kGreaterThanEqual));
+    EXPECT_TRUE(p->peek_is(Token::Type::kGreaterThanEqual))
+        << "expected: <= got: " << p->peek().to_name();
+    EXPECT_TRUE(p->peek_is(Token::Type::kVec2, 1))
+        << "expected: vec2 got: " << p->peek(1).to_name();
+    EXPECT_TRUE(p->peek_is(Token::Type::kLessThan, 2))
+        << "expected: < got: " << p->peek(2).to_name();
+}
+
+TEST_F(ParserImplTest, Peek_PastEnd) {
+    auto p = parser(">");
+    EXPECT_TRUE(p->peek_is(Token::Type::kGreaterThan));
+    EXPECT_TRUE(p->peek_is(Token::Type::kEOF, 1));
+    EXPECT_TRUE(p->peek_is(Token::Type::kEOF, 2));
+}
+
+TEST_F(ParserImplTest, Peek_PastEnd_WalkingPlaceholders) {
+    auto p = parser(">= >=");
+    auto& n = p->next();
+    ASSERT_TRUE(n.Is(Token::Type::kGreaterThanEqual));
+    EXPECT_TRUE(p->peek_is(Token::Type::kGreaterThanEqual))
+        << "expected: <= got: " << p->peek().to_name();
+    EXPECT_TRUE(p->peek_is(Token::Type::kEOF, 1)) << "expected: EOF got: " << p->peek(1).to_name();
+}
+
+TEST_F(ParserImplTest, Peek_AfterSplit) {
+    auto p = parser(">= vec2<u32>");
+    auto& n = p->next();
+    ASSERT_TRUE(n.Is(Token::Type::kGreaterThanEqual));
+    EXPECT_TRUE(p->peek_is(Token::Type::kVec2)) << "expected: vec2 got: " << p->peek().to_name();
+
+    p->split_token(Token::Type::kGreaterThan, Token::Type::kEqual);
+    ASSERT_TRUE(n.Is(Token::Type::kGreaterThan));
+    EXPECT_TRUE(p->peek_is(Token::Type::kEqual)) << "expected: = got: " << p->peek().to_name();
+}
+
 }  // namespace
 }  // namespace tint::reader::wgsl