[clangd] Avoid scanning up to end of file on each comment!

[sammccall]

Summary

Assigning char* (pointing at comment start) to StringRef was causing us
to scan the rest of the source file looking for the null terminator.

This seems to be eating about 8% of our *total* CPU!

While fixing this, factor out the common bits from the two places we're
parsing IWYU pragmas.

[kadircet]

wow, that's an interesting finding. thanks!

clang-tools-extra/clangd/Headers.cpp

25–35

i think this also deserves a comment to make sure people won't refactor it in the future to take in a stringref.

25–35

i'd actually return an empty stringref, instead of None. it makes logic in callers less complicated and I don't think we convey any different signal between None vs empty right now (at least a signal that can be used by the callers)

27

nit: static constexpr StringLiteral instead (we can just use size() afterwards instead of calling strlen a bunch.

clang-tools-extra/clangd/unittests/HeadersTests.cpp

456

it might be useful to get rid of the space between : and keep as well.

[sammccall]

clang-tools-extra/clangd/Headers.cpp

25–35

The comment is in the header file, extended it to be more explicit about avoiding StringRef

25–35

it makes logic in callers less complicated

I think this encourages a risky pattern for very minimal gain.
We've seen comment handlers are performance-sensitive (each strlen of the source file isn't that expensive!).
If we bail out explicitly early in the overwhelmingly common no-pragma case, then the remaining code is not critical.
If we continue on and "naturally" fail to match any pragma case, I don't see performance problems today, however we need to be extra-careful whenever we change the code. I'd much prefer to have to read/write an obvious if (empty) return than to think about performance.

If we're going to prefer an early bailout, StringRef isn't really simpler than Optional<StringRef>, and it pushes callers to handle this the right way.

[sammccall]

BTW it looks like the 8% profile I saw was an outlier, 3-4% seems more typical.

clang-tools-extra/clangd/Headers.h

 #include <string>
 
 namespace clang {
 namespace clangd {
 
 /// Returns true if \p Include is literal include like "path" or <path>.
 bool isLiteralInclude(llvm::StringRef Include);
 
+/// If Text begins an Include-What-You-Use directive, returns it.
+/// Given "// IWYU pragma: keep", returns "keep".
+/// Input is a null-terminated char* as provided by SM.getCharacterData().
+/// (This should not be StringRef as we do *not* want to scan for its length).
+llvm::Optional<StringRef> parseIWYUPragma(const char *Text);
+
 /// Represents a header file to be #include'd.
 struct HeaderFile {
   std::string File;
   /// If this is true, `File` is a literal string quoted with <> or "" that
   /// can be #included directly; otherwise, `File` is an absolute file path.
   bool Verbatim;
 
   bool valid() const;

clang-tools-extra/clangd/Headers.cpp

 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 #include <cstring>
 
 namespace clang {
 namespace clangd {
 
-const char IWYUPragmaKeep[] = "// IWYU pragma: keep";
-const char IWYUPragmaExport[] = "// IWYU pragma: export";
-const char IWYUPragmaBeginExports[] = "// IWYU pragma: begin_exports";
+llvm::Optional<StringRef> parseIWYUPragma(const char *Text) {
+  // This gets called for every comment seen in the preamble, so it's quite hot.
+  constexpr llvm::StringLiteral IWYUPragma = "// IWYU pragma: ";
+  if (strncmp(Text, IWYUPragma.data(), IWYUPragma.size()))
+    return llvm::None;
+  Text += IWYUPragma.size();
+  const char *End = Text;
+  while (*End != 0 && *End != '\n')
+    ++End;
+  return StringRef(Text, End - Text);
+}
 
 class IncludeStructure::RecordHeaders : public PPCallbacks,
                                         public CommentHandler {
 public:
   RecordHeaders(const CompilerInstance &CI, IncludeStructure *Out)
       : SM(CI.getSourceManager()),
         HeaderInfo(CI.getPreprocessor().getHeaderSearchInfo()), Out(Out) {}
 
void FileChanged(SourceLocation Loc, FileChangeReason Reason,
     }
     case PPCallbacks::RenameFile:
     case PPCallbacks::SystemHeaderPragma:
       break;
     }
   }
 
   bool HandleComment(Preprocessor &PP, SourceRange Range) override {
-    bool Err = false;
-    llvm::StringRef Text = SM.getCharacterData(Range.getBegin(), &Err);
-    if (Err)
+    auto Pragma = parseIWYUPragma(SM.getCharacterData(Range.getBegin()));
+    if (!Pragma)
       return false;
+
     if (inMainFile()) {
       // Given:
       //
       // #include "foo.h"
       // #include "bar.h" // IWYU pragma: keep
       //
       // The order in which the callbacks will be triggered:
       //
       // 1. InclusionDirective("foo.h")
       // 2. handleCommentInMainFile("// IWYU pragma: keep")
       // 3. InclusionDirective("bar.h")
       //
       // This code stores the last location of "IWYU pragma: keep" (or export)
       // comment in the main file, so that when InclusionDirective is called, it
       // will know that the next inclusion is behind the IWYU pragma.
       // FIXME: Support "IWYU pragma: begin_exports" and "IWYU pragma:
       // end_exports".
-      if (!Text.startswith(IWYUPragmaExport) &&
-          !Text.startswith(IWYUPragmaKeep))
+      if (!Pragma->startswith("export") && !Pragma->startswith("keep"))
         return false;
       unsigned Offset = SM.getFileOffset(Range.getBegin());
       LastPragmaKeepInMainFileLine =
           SM.getLineNumber(SM.getMainFileID(), Offset) - 1;
     } else {
       // Memorize headers that that have export pragmas in them. Include Cleaner
       // does not support them properly yet, so they will be not marked as
       // unused.
       // FIXME: Once IncludeCleaner supports export pragmas, remove this.
-      if (!Text.startswith(IWYUPragmaExport) &&
-          !Text.startswith(IWYUPragmaBeginExports))
+      if (!Pragma->startswith("export") && !Pragma->startswith("begin_exports"))
         return false;
       Out->HasIWYUExport.insert(
           *Out->getID(SM.getFileEntryForID(SM.getFileID(Range.getBegin()))));
     }
     return false;
   }
 
 private:

clang-tools-extra/clangd/index/CanonicalIncludes.cpp

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
 #include "llvm/Support/Path.h"
 #include <algorithm>
 
 namespace clang {
 namespace clangd {
 namespace {
-const char IWYUPragma[] = "// IWYU pragma: private, include ";
-
 const std::pair<llvm::StringRef, llvm::StringRef> IncludeMappings[] = {
     {"include/__stddef_max_align_t.h", "<cstddef>"},
     {"include/__wmmintrin_aes.h", "<wmmintrin.h>"},
     {"include/__wmmintrin_pclmul.h", "<wmmintrin.h>"},
     {"include/adxintrin.h", "<immintrin.h>"},
     {"include/ammintrin.h", "<ammintrin.h>"},
     {"include/avx2intrin.h", "<immintrin.h>"},
     {"include/avx512bwintrin.h", "<immintrin.h>"},

 
 std::unique_ptr<CommentHandler>
 collectIWYUHeaderMaps(CanonicalIncludes *Includes) {
   class PragmaCommentHandler : public clang::CommentHandler {
   public:
     PragmaCommentHandler(CanonicalIncludes *Includes) : Includes(Includes) {}
 
     bool HandleComment(Preprocessor &PP, SourceRange Range) override {
-      llvm::StringRef Text =
-          Lexer::getSourceText(CharSourceRange::getCharRange(Range),
-                               PP.getSourceManager(), PP.getLangOpts());
-      if (!Text.consume_front(IWYUPragma))
+      auto Pragma = parseIWYUPragma(
+          PP.getSourceManager().getCharacterData(Range.getBegin()));
+      if (!Pragma || !Pragma->consume_front("private, include "))
         return false;
       auto &SM = PP.getSourceManager();
       // We always insert using the spelling from the pragma.
       if (auto *FE = SM.getFileEntryForID(SM.getFileID(Range.getBegin())))
-        Includes->addMapping(
-            FE->getLastRef(),
-            isLiteralInclude(Text) ? Text.str() : ("\"" + Text + "\"").str());
+        Includes->addMapping(FE->getLastRef(),
+                             isLiteralInclude(*Pragma)
+                                 ? Pragma->str()
+                                 : ("\"" + *Pragma + "\"").str());
       return false;
     }
 
   private:
     CanonicalIncludes *const Includes;
   };
   return std::make_unique<PragmaCommentHandler>(Includes);
 }

clang-tools-extra/clangd/unittests/HeadersTests.cpp

 //===-- HeadersTests.cpp - Include headers unit tests -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Headers.h"
 
 #include "Compiler.h"
+#include "Matchers.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
 using ::testing::Contains;
 using ::testing::ElementsAre;
+using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::Not;
 using ::testing::UnorderedElementsAre;
 
 class HeadersTest : public ::testing::Test {
 public:
   HeadersTest() {
     CDB.ExtraClangFlags = {SearchDirArg.c_str()};

 )cpp";
 
   auto Includes = collectIncludes();
   EXPECT_TRUE(Includes.hasIWYUExport(getID("export.h", Includes)));
   EXPECT_TRUE(Includes.hasIWYUExport(getID("begin_exports.h", Includes)));
   EXPECT_FALSE(Includes.hasIWYUExport(getID("none.h", Includes)));
 }
 
+TEST(Headers, ParseIWYUPragma) {
+  EXPECT_THAT(parseIWYUPragma("// IWYU pragma: keep"), HasValue(Eq("keep")));
+  EXPECT_THAT(parseIWYUPragma("// IWYU pragma: keep\netc"),
+              HasValue(Eq("keep")));
+  EXPECT_EQ(parseIWYUPragma("/* IWYU pragma: keep"), llvm::None)
+      << "Only // comments supported!";
+  EXPECT_EQ(parseIWYUPragma("//  IWYU pragma: keep"), llvm::None)
+      << "Sensitive to whitespace";
+  EXPECT_EQ(parseIWYUPragma("// IWYU pragma:keep"), llvm::None)
+      << "Sensitive to whitespace";
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang