Record //@... directive lines as comments when lexing. (#7494)

The `//@include-in-dumps` and `//@dump-sem-ir-begin`/`-end` tooling
directives were consumed for their side effects without a comment
record, so the tokens and comments together no longer reconstructed the
source: tooling that re-emits a file from them, such as `carbon format`,
silently dropped the directive lines. Now each recognized directive line
is also recorded as an ordinary full-line comment alongside its side
effect.

Adjacent full-line comments coalesce into one comment record only within
a category, determined by the byte after the `//` introducer: ordinary
comments (whitespace, or the end of the line or file), `//@...`
directives, and invalid introducers. A transition between categories
starts a new record, so a directive next to a comment block is its own
comment, while all the invalid spellings lump together to keep the
diagnostic noise at one per run.

The category boundary also fixes a lost directive: the invalid-comment
bulk skip compared only the `//` prefix, so `//!x` directly above
`//@dump-sem-ir-begin` absorbed the directive line and its side effect
was never recorded. Invalid comment runs now skip line by line (they
start from a diagnosed error, so they are not hot) and stop at a
whitespace or `@` introducer; the SIMD bulk skip handles only ordinary
comment blocks, whose prefix comparison already includes the whitespace
byte.

Nothing outside the lexer and the formatter reads comment records, and
lex dumps do not include comments, so no other behavior changes.

Assisted-by: Claude Code

toolchain/lex/lex.cpp

@@ -964,7 +964,7 @@ auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
   // byte.
   const bool is_trailing = position != line_info.start + line_info.indent;
 
-  // The introducer '//' must be followed by whitespace or EOF.
+  // Check whether the `//` introducer is followed by something valid.
   bool is_valid_after_slashes = true;
   if (position + 2 < static_cast<ssize_t>(source_text.size()) &&
       LLVM_UNLIKELY(!IsSpace(source_text[position + 2]))) {
@@ -972,19 +972,28 @@ auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
     // The `//@...` directives are tooling markers that are only meaningful as
     // full-line comments, so we only recognize them when not trailing.
     if (!is_trailing) {
+      // A directive is also recorded as a comment: the tokens and comments
+      // together reconstruct the source, so tooling such as the formatter
+      // would otherwise silently drop the directive line.
+      auto add_directive_comment_line = [&] {
+        buffer_.AddComment(line_info.indent, comment_start,
+                           buffer_.line_infos_.Get(next_line()).start,
+                           /*is_trailing=*/false);
+        AdvanceToLine(source_text, position, next_line());
+      };
       if (comment_text.starts_with("//@include-in-dumps\n")) {
         buffer_.has_include_in_dumps_ = true;
-        AdvanceToLine(source_text, position, next_line());
+        add_directive_comment_line();
         return;
       }
       if (comment_text.starts_with("//@dump-sem-ir-begin\n")) {
         BeginDumpSemIRRange(comment_text.begin());
-        AdvanceToLine(source_text, position, next_line());
+        add_directive_comment_line();
         return;
       }
       if (comment_text.starts_with("//@dump-sem-ir-end\n")) {
         EndDumpSemIRRange(comment_text.begin());
-        AdvanceToLine(source_text, position, next_line());
+        add_directive_comment_line();
         return;
       }
     }
@@ -1017,13 +1026,16 @@ auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
   }
 
   // A very common pattern is a long block of comment lines all with the same
-  // indent and comment start. We skip these comment blocks in bulk both for
-  // speed and to reduce redundant diagnostics if each line has the same
-  // erroneous comment start like `//!`.
+  // indent and comment start. We skip these comment blocks in bulk for speed,
+  // and with SIMD support short indents can be scanned extremely quickly; we
+  // expect these to be the dominant cases.
   //
-  // When we have SIMD support this is even more important for speed, as short
-  // indents can be scanned extremely quickly with SIMD and we expect these to
-  // be the dominant cases.
+  // An invalid comment start was already diagnosed above, so its block is
+  // instead skipped line by line below: a run of invalid comment lines lumps
+  // into one block regardless of which invalid byte follows each `//`, keeping
+  // the diagnostic noise to one per run, while a line whose introducer is
+  // valid (whitespace, or a `//@...` directive) ends the run and is lexed on
+  // its own.
   //
   // TODO: We should extend this to 32-byte SIMD on platforms with support.
   constexpr int MaxIndent = 13;
@@ -1038,7 +1050,7 @@ auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
     next_line_info.indent = indent;
     position = next_line_info.start;
   };
-  if (CARBON_USE_SIMD &&
+  if (CARBON_USE_SIMD && is_valid_after_slashes &&
       position + 16 < static_cast<ssize_t>(source_text.size()) &&
       indent <= MaxIndent) {
     // Load a mask based on the amount of text we want to compare.
@@ -1085,15 +1097,57 @@ auto Lexer::LexComment(llvm::StringRef source_text, ssize_t& position) -> void {
 #else
 #error "Unsupported SIMD architecture!"
 #endif
-    // TODO: If we finish the loop due to the position approaching the end of
-    // the buffer we may fail to skip the last line in a comment block that
-    // has an invalid initial sequence and thus emit extra diagnostics. We
-    // should really fall through to the generic skipping logic, but the code
-    // organization will need to change significantly to allow that.
   } else {
-    while (position + prefix_size < static_cast<ssize_t>(source_text.size()) &&
-           memcmp(source_text.data() + first_line_start,
-                  source_text.data() + position, prefix_size) == 0) {
+    auto continues_block = [&](ssize_t position) -> bool {
+      // Make sure the source text extends far enough for us to continue the
+      // block.
+      if (position + prefix_size > static_cast<ssize_t>(source_text.size())) {
+        return false;
+      }
+
+      // Check that the prefix matches. Otherwise, the block is done.
+      if (memcmp(source_text.data() + first_line_start,
+                 source_text.data() + position, prefix_size) != 0) {
+        return false;
+      }
+
+      // For something valid after `//`, we're done as we've ensured it was the
+      // _same_ valid suffix in the `memcmp`.
+      if (LLVM_LIKELY(is_valid_after_slashes)) {
+        return true;
+      }
+
+      // Past here, the block is a run of invalid comment lines and the
+      // matched prefix only covers this line's `//`, so examine what follows
+      // those slashes: only another invalid comment line continues the block,
+      // while a valid comment or directive ends it and is lexed on its own.
+
+      // A `//` that ends the source is a valid comment, so it doesn't
+      // continue the block.
+      if (position + prefix_size == static_cast<ssize_t>(source_text.size())) {
+        return false;
+      }
+
+      char after_slashes = source_text[position + prefix_size];
+
+      // Whitespace after the `//` makes this line a valid comment, so it
+      // doesn't continue the block.
+      if (IsSpace(after_slashes)) {
+        return false;
+      }
+
+      // An `@` makes this line a `//@...` directive that must be lexed on its
+      // own to recognize its side effects, so it doesn't continue the block.
+      if (after_slashes == '@') {
+        return false;
+      }
+
+      // Anything else is another invalid comment line continuing the block.
+      return true;
+    };
+
+    // Skip lines that are combined into a comment block.
+    while (continues_block(position)) {
       skip_to_next_line();
     }
   }

toolchain/lex/tokenized_buffer.cpp

@@ -418,6 +418,33 @@ auto TokenizedBuffer::IsTrailingComment(CommentIndex comment_index) const
   return comments_.Get(comment_index).is_trailing;
 }
 
+namespace {
+// The category of a full-line comment, determined by the byte after the `//`
+// introducer. Adjacent full-line comments coalesce only within a category, so
+// a comment's category is well defined by its first line.
+enum class CommentCategory : uint8_t {
+  // Whitespace, or nothing before the end of the line or file.
+  Ordinary,
+  // A `//@...` tooling directive.
+  Directive,
+  // Every other byte; the invalid spellings are lumped into one category.
+  Invalid,
+};
+}  // namespace
+
+// Returns the comment's category; see `CommentCategory`.
+static auto GetCommentCategory(llvm::StringRef source, int32_t comment_start)
+    -> CommentCategory {
+  if (comment_start + 2 >= static_cast<int32_t>(source.size()) ||
+      IsSpace(source[comment_start + 2])) {
+    return CommentCategory::Ordinary;
+  }
+  if (source[comment_start + 2] == '@') {
+    return CommentCategory::Directive;
+  }
+  return CommentCategory::Invalid;
+}
+
 auto TokenizedBuffer::AddComment(int32_t indent, int32_t start, int32_t end,
                                  bool is_trailing) -> void {
   // A comment runs forward from its start, and its length is stored in 31 bits
@@ -425,13 +452,17 @@ auto TokenizedBuffer::AddComment(int32_t indent, int32_t start, int32_t end,
   // in 31 bits because the source size is bounded by `INT32_MAX`.
   CARBON_DCHECK(start <= end);
 
-  // A block of identical full-line comments is coalesced into a single comment.
-  // A trailing comment is always standalone: it never extends a preceding
-  // comment, nor is it extended by a following one.
+  // A block of adjacent full-line comments in the same category is coalesced
+  // into a single comment; transitioning between ordinary comments, `//@...`
+  // directives, and invalid introducers starts a new one. A trailing comment
+  // is always standalone: it never extends a preceding comment, nor is it
+  // extended by a following one.
   if (!is_trailing && comments_.size() > 0) {
     auto& comment = comments_.Get(CommentIndex(comments_.size() - 1));
     if (!comment.is_trailing &&
-        comment.start + comment.length + indent == start) {
+        comment.start + comment.length + indent == start &&
+        GetCommentCategory(source_->text(), comment.start) ==
+            GetCommentCategory(source_->text(), start)) {
       CARBON_DCHECK(comment.start <= end);
       comment.length = end - comment.start;
       return;

toolchain/lex/tokenized_buffer_test.cpp

@@ -1163,6 +1163,109 @@ TEST_F(LexerTest, TrailingCommentAfterMultiLineString) {
   EXPECT_TRUE(buffer.IsTrailingComment(CommentIndex(0)));
 }
 
+TEST_F(LexerTest, DirectiveComments) {
+  // A `//@...` directive line is consumed for its tooling side effects and is
+  // also recorded as a comment: the tokens and comments together reconstruct
+  // the source, so tooling such as the formatter preserves the directive.
+  auto& buffer = compile_helper_.GetTokenizedBuffer(
+      "//@include-in-dumps\n"
+      "\n"
+      "//@dump-sem-ir-begin\n"
+      "var x: i32 = 0;\n"
+      "//@dump-sem-ir-end\n");
+  EXPECT_FALSE(buffer.has_errors());
+  EXPECT_TRUE(buffer.has_include_in_dumps());
+  EXPECT_TRUE(buffer.has_dump_sem_ir_ranges());
+  ASSERT_THAT(buffer.comments_size(), Eq(3));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(0)),
+              Eq("//@include-in-dumps\n"));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(1)),
+              Eq("//@dump-sem-ir-begin\n"));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(2)),
+              Eq("//@dump-sem-ir-end\n"));
+  EXPECT_FALSE(buffer.IsTrailingComment(CommentIndex(0)));
+
+  // Adjacent full-line comments coalesce only within a category: a directive
+  // next to an ordinary comment is its own record, while adjacent directives
+  // share one.
+  auto& adjacent = compile_helper_.GetTokenizedBuffer(
+      "// Dump this file.\n"
+      "//@include-in-dumps\n"
+      "//@dump-sem-ir-begin\n"
+      "var x: i32 = 0;\n"
+      "//@dump-sem-ir-end\n");
+  EXPECT_FALSE(adjacent.has_errors());
+  EXPECT_TRUE(adjacent.has_include_in_dumps());
+  ASSERT_THAT(adjacent.comments_size(), Eq(3));
+  EXPECT_THAT(adjacent.GetCommentText(CommentIndex(0)),
+              Eq("// Dump this file.\n"));
+  EXPECT_THAT(adjacent.GetCommentText(CommentIndex(1)),
+              Eq("//@include-in-dumps\n//@dump-sem-ir-begin\n"));
+  EXPECT_THAT(adjacent.GetCommentText(CommentIndex(2)),
+              Eq("//@dump-sem-ir-end\n"));
+}
+
+TEST_F(LexerTest, DirectiveAfterInvalidComment) {
+  // An invalid comment line does not absorb a following directive: the
+  // directive ends the invalid run, so its side effects are still recognized
+  // and it is recorded separately.
+  auto& buffer = compile_helper_.GetTokenizedBuffer(
+      "//!bad\n"
+      "//@dump-sem-ir-begin\n"
+      "var x: i32 = 0;\n"
+      "//@dump-sem-ir-end\n");
+  EXPECT_TRUE(buffer.has_errors());
+  EXPECT_TRUE(buffer.has_dump_sem_ir_ranges());
+  ASSERT_THAT(buffer.comments_size(), Eq(3));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(0)), Eq("//!bad\n"));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(1)),
+              Eq("//@dump-sem-ir-begin\n"));
+}
+
+TEST_F(LexerTest, InvalidCommentRunsLumpTogether) {
+  // A run of invalid comment lines lumps into one comment and one diagnostic,
+  // no matter which invalid byte follows each `//`; a valid comment ends the
+  // run and starts its own record.
+  Testing::MockDiagnosticConsumer consumer;
+  EXPECT_CALL(consumer,
+              HandleDiagnostic(IsSingleDiagnostic(
+                  Diagnostics::Kind::NoWhitespaceAfterCommentIntroducer,
+                  Diagnostics::Level::Error, 1, 3, _)));
+  auto& buffer = compile_helper_.GetTokenizedBuffer(
+      "//!one\n"
+      "//?two\n"
+      "// valid\n",
+      &consumer);
+  ASSERT_THAT(buffer.comments_size(), Eq(2));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(0)), Eq("//!one\n//?two\n"));
+  EXPECT_THAT(buffer.GetCommentText(CommentIndex(1)), Eq("// valid\n"));
+}
+
+TEST_F(LexerTest, InvalidCommentRunAtEof) {
+  // An invalid run's line-by-line skip reads the byte after each line's `//`;
+  // these cases end the source at that read's boundary, with no trailing
+  // newline, to pin its bounds.
+
+  // The final line's introducer byte is the last byte of the source.
+  auto& tight = compile_helper_.GetTokenizedBuffer(
+      "//!a\n"
+      "//!");
+  EXPECT_TRUE(tight.has_errors());
+  ASSERT_THAT(tight.comments_size(), Eq(1));
+  EXPECT_THAT(tight.GetCommentText(CommentIndex(0)), Eq("//!a\n//!"));
+
+  // A bare `//` ending the source has no byte after the introducer at all:
+  // the run must stop before it rather than read past the end, and it lexes
+  // as its own valid, empty comment.
+  auto& bare = compile_helper_.GetTokenizedBuffer(
+      "//!a\n"
+      "//");
+  EXPECT_TRUE(bare.has_errors());
+  ASSERT_THAT(bare.comments_size(), Eq(2));
+  EXPECT_THAT(bare.GetCommentText(CommentIndex(0)), Eq("//!a\n"));
+  EXPECT_THAT(bare.GetCommentText(CommentIndex(1)), Eq("//"));
+}
+
 TEST_F(LexerTest, DiagnosticWhitespace) {
   Testing::MockDiagnosticConsumer consumer;
   EXPECT_CALL(consumer,
@@ -1318,14 +1421,14 @@ TEST_F(LexerTest, MultipleComments) {
 {4}
 x
 )";
-  constexpr llvm::StringLiteral Comments[] = {
+  constexpr llvm::StringLiteral Groups[] = {
       // NOLINTNEXTLINE(bugprone-suspicious-missing-comma)
       "// This comment should be possible to parse with SIMD.\n"
       "// This one too.\n",
       "// This one as well, though it's a different indent.\n"
       "        // And mixes indent.\n"
       "   // And mixes indent more.\n",
-      "// This is one comment:\n"
+      "// This is several comments:\n"
       "//Invalid\n"
       "// Valid\n"
       "//Invalid\n"
@@ -1334,18 +1437,35 @@ x
       "//\n"
       "// Valid\n",
       "// This uses a high indent, which stops SIMD.\n", "//\n"};
-  std::string source = llvm::formatv(Format, Comments[0], Comments[1],
-                                     Comments[2], Comments[3], Comments[4])
+  std::string source = llvm::formatv(Format, Groups[0], Groups[1], Groups[2],
+                                     Groups[3], Groups[4])
                            .str();
 
   auto& buffer = compile_helper_.GetTokenizedBuffer(source);
   EXPECT_TRUE(buffer.has_errors());
 
-  EXPECT_THAT(buffer.comments_size(), Eq(std::size(Comments)));
-  for (int i :
-       llvm::seq(std::min<int>(buffer.comments_size(), std::size(Comments)))) {
+  // The third group splits at each transition between ordinary and invalid
+  // comment introducers; the other groups each form one comment.
+  constexpr llvm::StringLiteral ExpectedComments[] = {
+      // NOLINTNEXTLINE(bugprone-suspicious-missing-comma)
+      "// This comment should be possible to parse with SIMD.\n"
+      "// This one too.\n",
+      "// This one as well, though it's a different indent.\n"
+      "        // And mixes indent.\n"
+      "   // And mixes indent more.\n",
+      "// This is several comments:\n", "//Invalid\n", "// Valid\n",
+      "//Invalid\n",
+      // NOLINTNEXTLINE(bugprone-suspicious-missing-comma)
+      "//\n"
+      "// Valid\n"
+      "//\n"
+      "// Valid\n",
+      "// This uses a high indent, which stops SIMD.\n", "//\n"};
+  EXPECT_THAT(buffer.comments_size(), Eq(std::size(ExpectedComments)));
+  for (int i : llvm::seq(std::min<int>(buffer.comments_size(),
+                                       std::size(ExpectedComments)))) {
     EXPECT_THAT(buffer.GetCommentText(CommentIndex(i)).str(),
-                testing::StrEq(Comments[i]));
+                testing::StrEq(ExpectedComments[i]));
   }
   EXPECT_THAT(buffer, HasTokens(llvm::ArrayRef<ExpectedToken>{
                           {.kind = TokenKind::FileStart},