[clangd] Record IWYU pragma keep in the IncludeStructure

[kbobyrev]

Summary

This will allow the IncludeCleaner to suppress warnings on the lines with "IWYU
pragma: keep".

Clang APIs are not very convinient, so the code has to navigate around it.

[kadircet]

It's annoying that we see comments and inclusion directives out-of-order, we can try fixing it on the parser side (I think it is incidental that these are issued in that order currently, they are eagerly trying to generate a fix/diagnostic for tokens after a pp-directive. hence we can first issue the PP callback, and then diagnose for rest of the line instead).
But I don't think it matters in the long run especially when we want to handle different types of pragmas in the future (which is the plan IIRC), they might not even necessarily be in the main file, let alone being seen with right sequencing.

So I'd suggest just building a side table information about these interesting comments/pragmas we see as we go, and then merge them with rest of the information we've built inside RecordHeaders::EndOfMainFile.
That'll enable us to separate concerns here. Instead of piggy backing CommentHandler on RecordHeaders, we can have a separate comment handler exposed via IncludeStructure::commentHandler() (similar to collect) and let it update this side table of pragma/comment information.

WDYT?

[sammccall]

Sorry for forgetting about this one. Hopefully I can still help now by disagreeing with Kadir and creating an awkward stalemate instead.

I agree that the ownership situation is ugly, but I think a single object here is still simpler. Partly because the ordering isn't some weird coincidence (see below)

clang-tools-extra/clangd/Headers.cpp

111

I don't think this has anything to do with ownership or registration order, rather events get recorded in order that constructs *finish* rather than start.

The PP doesn't call InclusionDirective until it sees the end of the line, and it has to skip over the comment to get there.

(I like the example and the sequencing, I'd just replace the first paragraph with "Given:" or so)

clang-tools-extra/clangd/Headers.h

132

Alternatively we could just give up on having a pure-ish API and say void collect(Preprocessor&).
Again there's not really any other sensible way to use it.

[sammccall]

Didn't really address this part

But I don't think it matters in the long run especially when we want to handle different types of pragmas in the future
So I'd suggest just building a side table information

That's a more general approach for sure. But it does seem there's some complexity: what is the data structure and how do we merge it back into the graph. (e.g. are we going to record line numbers of each #include in the graph for this merging?)

So let's think about the concrete cases:

• export
• begin_exports...end_exports
• private maybe?

The TL;DR is that each of these either stand alone, or we see the directive before the inclusion it describes in a way that's easy to keep track of incrementally. So I'm not really seeing much implementation difficulty to justify trying to move the directives and inclusion processing apart.

Export seems basically similar. We're going to see the pragma immediately before the directive.
The same stateful solution seems just as clear, I don't see any problem with it not being in the main file.
(There's the "same line" issue of having to look up line tables for everything, maybe we can dodge it by comparing SourceLocations in raw form instead, or use the statefulness to avoid checks in most cases)

Begin/end is a bit different. The stateful solution looks like: keep a Set<FileID> where exports are turned on, mutate it when you see a directive, and check it when you seen an inclusion.
The side-table solution is a bit more complicated: a map of fileIDs to a list of ranges or something. Manageable for sure.

Private I think we have to treat as a standalone directive that mutates the current file, and the reference to another file is treated as a string. It's not a file the PP has necessarily seen. So the result is a map<file, string> and it seems adequate for this to be written straight into IncludeStructure.

clang-tools-extra/clangd/Headers.cpp

127

This code will run for every comment in the preamble, so we should probably be careful about its performance.

I'd guess the cheapest thing is:

• check that Range.begin() is not a macro
• SM.getCharacterData(Range.begin()) yields a char* that's guaranteed null-terminated
• try to consume the IWYU pragma: prefix

This doesn't deal with makeFileCharRange, touch End at all, etc.

[kadircet]

Sorry for forgetting about this one. Hopefully I can still help now by disagreeing with Kadir and creating an awkward stalemate instead.

Haha, always welcome!

Partly because the ordering isn't some weird coincidence (see below)

Right, I suppose it makes sense when you look at the "construct being finalized" perspective.

That's a more general approach for sure. But it does seem there's some complexity: what is the data structure and how do we merge it back into the graph. (e.g. are we going to record line numbers of each #include in the graph for this merging?)

Definitely, but we are not going to get rid of that extra complexity no matter which solution we go with. We still need to store some extra state about these comments, the only difference is when we process them & whether we store just the latest or the whole set.
Surely processing them in a different callback then InclusionDirective would add extra complexity, but I think it might also be simpler in general since we don't need to think about all the sequencing of these operations.
But after looking at the particular cases, I think I agree with your suggestion overall. Relying on the current sequencing to keep the state incremental rather than whole history sounds like the better trade off here. (details below).

So let's think about the concrete cases:

• export
• begin_exports...end_exports
• private maybe?

The TL;DR is that each of these either stand alone, or we see the directive before the inclusion it describes in a way that's easy to keep track of incrementally. So I'm not really seeing much implementation difficulty to justify trying to move the directives and inclusion processing apart.

Yeah these are the cases I had in mind, but didn't really think about them thoroughly.
I guess in all of the export cases we just need to operate with the information we've seen before a particular inclusion directive, and the file including it. So there doesn't seem to be any need for postponing the processing. (Plus as you pointed out, keeping just the current state is definitely more simpler than keeping a set of those).

Private I think we have to treat as a standalone directive that mutates the current file, and the reference to another file is treated as a string. It's not a file the PP has necessarily seen. So the result is a map<file, string> and it seems adequate for this to be written straight into IncludeStructure.

Yup, I guess this one looks pretty much the same for both approaches.

[kbobyrev]

Change the API, update the patch. It's crashing ATM, investigating the
problems.

[kbobyrev]

Fix failures, address review comments.

clang-tools-extra/clangd/Headers.h

132

PrecompiledPreamble.cpp requires PPCallbacks being returned directly, so I think we can't easily have `void IncludeStructure::collect(PP &).

[kbobyrev]

gentle ping

[kadircet]

clang-tools-extra/clangd/Headers.cpp

126

we're already keeping track of being inside builtinfile today, through FileChanged callbacks, can you also keep track of being inside main file rather than doing the check for each comment we hit?

149

maybe LastPragmaKeepInMainFileLine also mention in the comments that this is 0-indexed?

clang-tools-extra/clangd/Headers.h

132

as discussed offline let's do that inside BeforeExecute callback.

[kbobyrev]

Resolve almost all comments.

[kbobyrev]

Address the last comment.

[kadircet]

thanks, LGTM! sorry for taking so long.

clang-tools-extra/clangd/Headers.cpp

139

maybe just Keeps track of include depth for the current file. It's 1 for main file. ?

202

s/move/std::move/g

clang-tools-extra/clangd/Headers.h

129–130

I suppose the comment needs updating too. (it no longer returns nor requires extra comment handler insertion).

Inserts a PPCallback and CommentHandler that'll populate this structure. ?

[kbobyrev]

Resolve review comments.

clang-tools-extra/clangd/CodeComplete.cpp

bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   }
   // Macros can be defined within the preamble region of the main file.
   // They don't fall nicely into our index/Sema dichotomy:
   //  - they're not indexed for completion (they're not available across files)
   //  - but Sema code complete won't see them: as part of the preamble, they're
   //    deserialized only when mentioned.
   // Force them to be deserialized so SemaCodeComplete sees them.
   loadMainFilePreambleMacros(Clang->getPreprocessor(), Input.Preamble);
-  if (Includes)
-    Clang->getPreprocessor().addPPCallbacks(Includes->collect(*Clang));
+  if (Includes) {
+    auto Collector = Includes->collect(*Clang);
+    Clang->getPreprocessor().addCommentHandler(Collector.get());
+    Clang->getPreprocessor().addPPCallbacks(move(Collector));
+  }
   if (llvm::Error Err = Action.Execute()) {
     log("Execute() failed when running codeComplete for {0}: {1}",
         Input.FileName, toString(std::move(Err)));
     return false;
   }
   Action.EndSourceFile();
 
   return true;

clang-tools-extra/clangd/Headers.h

 #include "support/Logger.h"
 #include "support/Path.h"
 #include "clang/Basic/FileEntry.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
+#include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
 #include "llvm/Support/VirtualFileSystem.h"

 struct Inclusion {
   tok::PPKeywordKind Directive; // Directive used for inclusion, e.g. import
   std::string Written;          // Inclusion name as written e.g. <vector>.
   Path Resolved; // Resolved path of included file. Empty if not resolved.
   unsigned HashOffset = 0; // Byte offset from start of file to #.
   int HashLine = 0;        // Line number containing the directive, 0-indexed.
   SrcMgr::CharacteristicKind FileKind = SrcMgr::C_User;
   llvm::Optional<unsigned> HeaderID;
+  bool BehindPragmaKeep = false; // Has IWYU pragma: keep right after.
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Inclusion &);
 bool operator==(const Inclusion &LHS, const Inclusion &RHS);
 
 // Contains information about one file in the build graph and its direct
 // dependencies. Doesn't own the strings it references (IncludeGraph is
 // self-contained).
 struct IncludeGraphNode {

 // in any non-preamble inclusions.
 class IncludeStructure {
 public:
   IncludeStructure() {
     // Reserve HeaderID = 0 for the main file.
     RealPathNames.emplace_back();
   }
 
+  class IncludeCollector : public PPCallbacks, public CommentHandler {};
   // Returns a PPCallback that visits all inclusions in the main file and
-  // populates the structure.
-  std::unique_ptr<PPCallbacks> collect(const CompilerInstance &CI);
+  // populates the structure. IncludeCollector can also scan the comments for
+  // IWYU pragmas but it needs to be explicitly added to as a preprocessor
+  // comment handler. PP wants to own the PPCallbacks, so the typical way to
+  // do both is:
+  //
+  // auto Collector = collect(SM);
+  // PP.addCommentHandler(Colletor.get());
+  // PP.addPPCallbacks(move(Collector));
+  //
+  // When IncludeCollector is not added as the CommentHandler, resulting
+  // include structure won't have the IWYU pragmas information.
+  std::unique_ptr<IncludeStructure::IncludeCollector>
+  collect(const CompilerInstance &CI);
 
   // HeaderID identifies file in the include graph. It corresponds to a
   // FileEntry rather than a FileID, but stays stable across preamble & main
   // file builds.
   enum class HeaderID : unsigned {};
 
   llvm::Optional<HeaderID> getID(const FileEntry *Entry) const;
   HeaderID getOrCreateID(const FileEntry *Entry);

clang-tools-extra/clangd/Headers.cpp

 #include "SourceCode.h"
 #include "support/Logger.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/HeaderSearch.h"
+#include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 
-class IncludeStructure::RecordHeaders : public PPCallbacks {
+const char IWYUPragmaKeep[] = "// IWYU pragma: keep";
+
+class IncludeStructure::RecordHeaders
+    : public IncludeStructure::IncludeCollector {
 public:
-  RecordHeaders(const SourceManager &SM, HeaderSearch &HeaderInfo,
-                IncludeStructure *Out)
-      : SM(SM), HeaderInfo(HeaderInfo), Out(Out) {}
+  RecordHeaders(const CompilerInstance &CI, IncludeStructure *Out)
+      : SM(CI.getSourceManager()),
+        HeaderInfo(CI.getPreprocessor().getHeaderSearchInfo()), Out(Out) {}
 
   // Record existing #includes - both written and resolved paths. Only #includes
   // in the main file are collected.
   void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                           llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange /*FilenameRange*/,
                           const FileEntry *File, llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
if (isInsideMainFile(HashLoc, SM)) {
       Inc.Resolved = std::string(File ? File->tryGetRealPathName() : "");
       Inc.HashOffset = SM.getFileOffset(HashLoc);
       Inc.HashLine =
           SM.getLineNumber(SM.getFileID(HashLoc), Inc.HashOffset) - 1;
       Inc.FileKind = FileKind;
       Inc.Directive = IncludeTok.getIdentifierInfo()->getPPKeywordID();
       if (File)
         Inc.HeaderID = static_cast<unsigned>(Out->getOrCreateID(File));
+      if (PragmaKeepInMainFile == Inc.HashLine)
+        Inc.BehindPragmaKeep = true;
     }
 
     // Record include graph (not just for main-file includes)
     if (File) {
       auto *IncludingFileEntry = SM.getFileEntryForID(SM.getFileID(HashLoc));
       if (!IncludingFileEntry) {
         assert(SM.getBufferName(HashLoc).startswith("<") &&
                "Expected #include location to be a file or <built-in>");
case PPCallbacks::ExitFile: {
       break;
     }
     case PPCallbacks::RenameFile:
     case PPCallbacks::SystemHeaderPragma:
       break;
     }
   }
 
+  // Given:
+  //
+  // #include "foo.h"
+  // #include "bar.h" // IWYU pragma: keep
+  //
+  // The order in which the callbacks will be triggered:
+  //
+  // 1. InclusionDirective("foo.h")
+  // 2. HandleComment("// IWYU pragma: keep")
+  // 3. InclusionDirective("bar.h")
+  //
+  // HandleComment will store the last location of "IWYU pragma: keep" comment
+  // in the main file, so that when InclusionDirective is called, it will know
+  // that the next inclusion is behind the IWYU pragma.
+  bool HandleComment(Preprocessor &PP, SourceRange Range) override {
+    if (!isInsideMainFile(Range.getBegin(), SM) || Range.getBegin().isMacroID())
+      return false;
+    bool Err = false;
+    llvm::StringRef Text = SM.getCharacterData(Range.getBegin(), &Err);
+    if (Err || !Text.consume_front(IWYUPragmaKeep))
+      return false;
+    unsigned Offset = SM.getFileOffset(Range.getBegin());
+    PragmaKeepInMainFile =
+        SM.getLineNumber(SM.getFileID(Range.getBegin()), Offset) - 1;
+    return false;
+  }
+
 private:
   const SourceManager &SM;
   HeaderSearch &HeaderInfo;
   // Set after entering the <built-in> file.
   FileID BuiltinFile;
   // Indicates whether <built-in> file is part of include stack.
   bool InBuiltinFile = false;
 
   IncludeStructure *Out;
+
+  // The last line "IWYU pragma: keep" was seen in the main file.
+  int PragmaKeepInMainFile = -1;
 };
 
 bool isLiteralInclude(llvm::StringRef Include) {
   return Include.startswith("<") || Include.startswith("\"");
 }
 
 bool HeaderFile::valid() const {
   return (Verbatim && isLiteralInclude(File)) ||
llvm::sort(Includes, [](const Symbol::IncludeHeaderWithReferences &LHS,
     return LHS.References > RHS.References;
   });
   llvm::SmallVector<llvm::StringRef, 1> Headers;
   for (const auto &Include : Includes)
     Headers.push_back(Include.IncludeHeader);
   return Headers;
 }
 
-std::unique_ptr<PPCallbacks>
+std::unique_ptr<IncludeStructure::IncludeCollector>
 IncludeStructure::collect(const CompilerInstance &CI) {
   auto &SM = CI.getSourceManager();
   MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
-  return std::make_unique<RecordHeaders>(
-      SM, CI.getPreprocessor().getHeaderSearchInfo(), this);
+  return std::make_unique<RecordHeaders>(CI, this);
 }
 
 llvm::Optional<IncludeStructure::HeaderID>
 IncludeStructure::getID(const FileEntry *Entry) const {
   // HeaderID of the main file is always 0;
   if (Entry == MainFileEntry) {
     return static_cast<IncludeStructure::HeaderID>(0u);
   }
   auto It = UIDToIndex.find(Entry->getUniqueID());
   if (It == UIDToIndex.end())

clang-tools-extra/clangd/ParsedAST.cpp

if (Preamble) {
     Includes.MainFileIncludes = Patch->preambleIncludes();
     // Replay the preamble includes so that clang-tidy checks can see them.
     ReplayPreamble::attach(Patch->preambleIncludes(), *Clang,
                            Patch->modifiedBounds());
   }
   // Important: collectIncludeStructure is registered *after* ReplayPreamble!
   // Otherwise we would collect the replayed includes again...
   // (We can't *just* use the replayed includes, they don't have Resolved path).
-  Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang));
+  auto Collector = Includes.collect(*Clang);
+  Clang->getPreprocessor().addCommentHandler(Collector.get());
+  Clang->getPreprocessor().addPPCallbacks(move(Collector));
   // Copy over the macros in the preamble region of the main file, and combine
   // with non-preamble macros below.
   MainFileMacros Macros;
   if (Preamble)
     Macros = Preamble->Macros;
   Clang->getPreprocessor().addPPCallbacks(
       std::make_unique<CollectMainFileMacros>(Clang->getSourceManager(),
                                               Macros));

clang-tools-extra/clangd/Preamble.cpp

if (Clang->getFrontendOpts().Inputs.empty())
     return error("compiler instance had no inputs");
   // We are only interested in main file includes.
   Clang->getPreprocessorOpts().SingleFileParseMode = true;
   PreprocessOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     return error("failed BeginSourceFile");
   Preprocessor &PP = Clang->getPreprocessor();
   IncludeStructure Includes;
-  PP.addPPCallbacks(Includes.collect(*Clang));
+  auto Collector = Includes.collect(*Clang);
+  PP.addCommentHandler(Collector.get());
+  PP.addPPCallbacks(move(Collector));
   ScannedPreamble SP;
   SP.Bounds = Bounds;
   PP.addPPCallbacks(
       std::make_unique<DirectiveCollector>(PP, SP.TextualDirectives));
   if (llvm::Error Err = Action.Execute())
     return std::move(Err);
   Action.EndSourceFile();
   SP.Includes = std::move(Includes.MainFileIncludes);

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

protected:
   }
 
   IncludeStructure collectIncludes() {
     Clang = setupClang();
     PreprocessOnlyAction Action;
     EXPECT_TRUE(
         Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]));
     IncludeStructure Includes;
-    Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang));
+    auto Collector = Includes.collect(*Clang);
+    Clang->getPreprocessor().addCommentHandler(Collector.get());
+    Clang->getPreprocessor().addPPCallbacks(move(Collector));
     EXPECT_FALSE(Action.Execute());
     Action.EndSourceFile();
     return Includes;
   }
 
   // Calculates the include path, or returns "" on error or header should not be
   // inserted.
   std::string calculate(PathRef Original, PathRef Preferred = "",
protected:
   IgnoringDiagConsumer IgnoreDiags;
   std::unique_ptr<CompilerInstance> Clang;
 };
 
 MATCHER_P(Written, Name, "") { return arg.Written == Name; }
 MATCHER_P(Resolved, Name, "") { return arg.Resolved == Name; }
 MATCHER_P(IncludeLine, N, "") { return arg.HashLine == N; }
 MATCHER_P(Directive, D, "") { return arg.Directive == D; }
+MATCHER_P(HasPragmaKeep, H, "") { return arg.BehindPragmaKeep == H; }
 
 MATCHER_P2(Distance, File, D, "") {
   if (arg.getFirst() != File)
     *result_listener << "file =" << static_cast<unsigned>(arg.getFirst());
   if (arg.getSecond() != D)
     *result_listener << "distance =" << arg.getSecond();
   return arg.getFirst() == File && arg.getSecond() == D;
 }
)cpp";
   // ms-compatibility changes meaning of #import, make sure it is turned off.
   CDB.ExtraClangFlags.push_back("-fno-ms-compatibility");
   EXPECT_THAT(collectIncludes().MainFileIncludes,
               UnorderedElementsAre(Directive(tok::pp_include),
                                    Directive(tok::pp_import),
                                    Directive(tok::pp_include_next)));
 }
 
+TEST_F(HeadersTest, IWYUPragmaKeep) {
+  FS.Files[MainFile] = R"cpp(
+#include "bar.h" // IWYU pragma: keep
+#include "foo.h"
+)cpp";
+
+  EXPECT_THAT(
+      collectIncludes().MainFileIncludes,
+      UnorderedElementsAre(AllOf(Written("\"foo.h\""), HasPragmaKeep(false)),
+                           AllOf(Written("\"bar.h\""), HasPragmaKeep(true))));
+}
+
 TEST_F(HeadersTest, InsertInclude) {
   std::string Path = testPath("sub/bar.h");
   FS.Files[Path] = "";
   EXPECT_EQ(calculate(Path), "\"bar.h\"");
 }
 
 TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
   MainFile = testPath("main.h");