[include-cleaner] Add PragmaIncludes which handles include-mapping pragmas.

[hokein]

Summary

PragmaIncludes captures the progma-based header-mapping information, it is used
in the "Location => Header" step to determine the final spelling header for a
symbol (rather than the header directive).

[hokein]

Post it early in order to get initial feedback on the APIs.It is not finished, I plan to implement the IWYU keep pragma and add unittest.

[sammccall]

(It's a bit hard to evaluate this without the other pieces, but I suppose we need to start somewhere)

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Hooks.h

38 ↗

(On Diff #468176)

IWYU pragmas requires a PP listener, we'll also need PP events for other purposes (determining which files are self-contained, finding #includes in the main file). Should these be the same listener and write into the same class?

The argument for fewer components is usability: the embedder (standalone tool, clangd, clang-tidy) is responsible for connecting these hooks to the compiler and feeding the results back into the analysis.

So the library is somewhat simpler to use if the number of parts is small.
We need (I think) at least one for the AST and one for the preprocessor, because they need to be connected at different parts of the lifecycle.

How much do we gain by splitting up further than that?

44 ↗

(On Diff #468176)

At first glance, this signature doesn't look right: keep applies to directives, not to files.
e.g.

#include "foo.h" // keep
#include "foo.h"

the second should be removed

(it's a silly example, but in the past we've confused ourselves into doing IO or complicated string matching code for things that are about spelling).

I'd suggest some representation like the line number (where the # appears, not where the comment is), though the details probably matter less than the concept.

47 ↗

(On Diff #468176)

FileEntry* rather than FileEntryRef - we don't have different mappings when we open the same file using different names.

(This seems like a nit here but Header will be something like variant<stdlib::Header, FileEntry*, ...> and inconsistency would be confusing)

53 ↗

(On Diff #468176)

Why are these UniqueID rather than FileEntry*? Is the idea to try to use this in clangd?

We hadn't planned on doing this, because pushing clangd's weird requirements around not using FileEntry's to the library, ability to clone the data structure etc seemed pretty intrusive. Since we switched from name-based to UniqueID it might be doable...

clang-tools-extra/include-cleaner/lib/Hooks.cpp

24 ↗

(On Diff #468176)

while we're adding this, we should probably support /* as well as //, it's allowed per spec (I didn't add it to clangd just to keep NFC)

42 ↗

(On Diff #468176)

nit: there's no need to check Invalid, getCharacterData always returns a valid pointer and it will never have a spurious IWYU pragma

[sammccall]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Hooks.h

38 ↗

(On Diff #468176)

After discussing offline: I think we expect this to be reused in clangd so separating it from the bits that won't is useful.
I don't think the scope is exactly clear though: which of these should we bundle together?

• IWYU pragmas related to headers (private, export)
• IWYU pragmas related to the main file (keep/no_include)
• other equivalent things like #pragma include_instead
• tracking whether headers are self-contained
• recording #include graph (not sure IWYU even needs this)
• tracking macro occurrences in the main file
• details of #include directives from the main file

I'm tempted to say 1/2/3/4 are in-scope (since they're all gathered from roughly the same source and needed for the same thing) and so PragmaIncludes seems a little too narrow. Suggest HeaderQuirks maybe?

53 ↗

(On Diff #468176)

After discussing offline, using UniqueID to ensure the results are usable across FileManagers seems reasonable, but I think it needs to be documented.

[hokein]

• rebase
• address comments
• implement shouldKeep
• add unittests

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Hooks.h

38 ↗

(On Diff #468176)

My original plan was 1/2/3 (4 probably yes, as they mostly share the same thing).

As a personal preference, I like the name PragmaIncludes as it perfectly matches 1/2/3, but not 4; while the HeaderQuirks is probably a better umbrella, but it seems not quite self-contained. for now, I'd keep it as-is, but happy to change if you have a strong preference.

44 ↗

(On Diff #468176)

right, good point!

53 ↗

(On Diff #468176)

Added comments.

clang-tools-extra/include-cleaner/lib/Hooks.cpp

24 ↗

(On Diff #468176)

added handling of /*

[kadircet]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

35

triple slashes (all over the header)

41

i think this interface is OK, if we're planning to implement another adapter on top of this class and use this only as a repository. Then have the adapter provide a more convenient interface/data structures for types of queries we'll have during Location => Header mapping inside include cleaner.

As those will most likely look like getHeaders(SourceLocation)/getHeaders(string) -> vector<pair<Header, Signals>>, without really caring about the details an implementation detail being mapped to some other file, or getting exports of a certain header (as discussed this is the piece that would require a little bit more detailed design).

49

it feels a little bit unfortunate from layering perspective that this struct also needs to somehow have specialized knowledge about the main file, but also having the same comment handler logic both in here and inside a main file specific handler seems like too much duplication. So let's go with it in this way for now, but keep this quirk in mind so that we can decide on a better place in the future.

54

why do we need to expose this?

58

we should also keep headers that're marked as export.

clang-tools-extra/include-cleaner/lib/Record.cpp

64

nit: we can keep track of the current FileID using FileChanged/LexedFileChanged events (also gets rid of the isWrittenInMainFile call above/below).

[sammccall]

LG from my side

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

58

export isn't in this patch.
I don't think the details of where this data comes belongs in the header, it's not important to either the data structure or the interface.

67

nit: naming a member after the keys rather than the values is confusing, consider IWYUPublic or just PublicMapping?

clang-tools-extra/include-cleaner/lib/Record.cpp

112

nit: we can inline these trivial implementations into the header

[hokein]

address comments

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

41

yeah, the plan is that this class is a low-level piece, only providing the data extracted from pragmas, and we will build another layer on top of this class (this is my next step).

54

no needed, move it to private.

67

rename to IWYUPublic.

[sammccall]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

41

My suggestion would be to see what it looks like coding directly against this first - it might be fine, or it might make it obvious that making this class a little higher level is a better choice.
Trying to place an abstraction around this kind of "bits and pieces" information is easy to get wrong ime

[kadircet]

thanks, LG from my side as well, just some small tweaks for tests.

clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

124

it'd be better to test with an include-directive that doesn't have a keep pragma attached

135

it'd be better to map this to a header that isn't part of the TU, to make sure we don't rely on a file entry being around for a particular spelling.

[hokein]

address comments around tests.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

 // These are the simplest way to connect include-cleaner logic to the parser,
 // but other ways are possible (for example clangd records includes separately).
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_INCLUDE_CLEANER_RECORD_H
 #define CLANG_INCLUDE_CLEANER_RECORD_H
 
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/FileSystem/UniqueID.h"
 #include <memory>
 #include <vector>
 
 namespace clang {
 class ASTConsumer;
 class ASTContext;
+class CompilerInstance;
 class Decl;
+class FileEntry;
+
 namespace include_cleaner {
 
+/// Captures #include mapping information. It analyses IWYU Pragma comments and
+/// other use-instead-like mechanisms (#pragma include_instead) on included
+/// files.
+///
+/// This is a low-level piece being used in the "Location => Header" analysis
+/// step to determine the final public header rather than the header directly
+/// defines the symbol.
+class PragmaIncludes {
+public:
+  /// Installs an analysing PPCallback and CommentHandler and populates results
+  /// to the structure.
+  void record(const CompilerInstance &CI);
+
+  /// Returns true if the given #include of the main-file should never be
+  /// removed.
+  bool shouldKeep(unsigned HashLineNumber) const {
+    return ShouldKeep.find(HashLineNumber) != ShouldKeep.end();
+  }
+
+  /// Returns the public mapping include for the given physical header file.
+  /// Returns "" if there is none.
+  llvm::StringRef getPublic(const FileEntry *File) const;
+
+private:
+  class RecordPragma;
+  /// 1-based Line numbers for the #include directives of the main file that
+  /// should always keep (e.g. has the `IWYU pragma: keep` or `IWYU pragma:
+  /// export` right after).
+  llvm::DenseSet</*LineNumber*/ unsigned> ShouldKeep;
+
+  /// The public header mapping by the IWYU private pragma.
+  //
+  // !!NOTE: instead of using a FileEntry* to identify the physical file, we
+  // deliberately use the UniqueID to ensure the result is stable across
+  // FileManagers (for clangd's preamble and main-file builds).
+  llvm::DenseMap<llvm::sys::fs::UniqueID, std::string /*VerbatimSpelling*/>
+      IWYUPublic;
+
+  // FIXME: add other IWYU supports (export etc)
+  // FIXME: add support for clang use_instead pragma
+  // FIXME: add selfcontained file.
+};
+
 // Contains recorded parser events relevant to include-cleaner.
 struct RecordedAST {
   // The consumer (when installed into clang) tracks declarations in this.
   std::unique_ptr<ASTConsumer> record();
 
   ASTContext *Ctx = nullptr;
   // The set of declarations written at file scope inside the main file.
   //

clang-tools-extra/include-cleaner/lib/Record.cpp

 //===--- Record.cpp - Record compiler events ------------------------------===//
 //
 // 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 "clang-include-cleaner/Record.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclGroup.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Lex/PPCallbacks.h"
+#include "clang/Lex/Preprocessor.h"
 
 namespace clang::include_cleaner {
 
+// FIXME: this is a mirror of clang::clangd::parseIWYUPragma, move to libTooling
+// to share the code?
+static llvm::Optional<StringRef> parseIWYUPragma(const char *Text) {
+  assert(strncmp(Text, "//", 2) || strncmp(Text, "/*", 2));
+  constexpr llvm::StringLiteral IWYUPragma = " IWYU pragma: ";
+  Text += 2; // Skip the comment start, // or /*.
+  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 PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
+public:
+  RecordPragma(const CompilerInstance &CI, PragmaIncludes *Out)
+      : SM(CI.getSourceManager()), Out(Out) {}
+
+  void FileChanged(SourceLocation Loc, FileChangeReason Reason,
+                   SrcMgr::CharacteristicKind FileType,
+                   FileID PrevFID) override {
+    InMainFile = SM.isWrittenInMainFile(Loc);
+  }
+
+  void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
+                          llvm::StringRef FileName, bool IsAngled,
+                          CharSourceRange /*FilenameRange*/,
+                          Optional<FileEntryRef> File,
+                          llvm::StringRef /*SearchPath*/,
+                          llvm::StringRef /*RelativePath*/,
+                          const clang::Module * /*Imported*/,
+                          SrcMgr::CharacteristicKind FileKind) override {
+    if (!InMainFile)
+      return;
+    int HashLine =
+        SM.getLineNumber(SM.getMainFileID(), SM.getFileOffset(HashLoc));
+    if (LastPragmaKeepInMainFileLine == HashLine)
+      Out->ShouldKeep.insert(HashLine);
+  }
+
+  bool HandleComment(Preprocessor &PP, SourceRange Range) override {
+    auto &SM = PP.getSourceManager();
+    auto Pragma = parseIWYUPragma(SM.getCharacterData(Range.getBegin()));
+    if (!Pragma)
+      return false;
+
+    if (Pragma->consume_front("private, include ")) {
+      // We always insert using the spelling from the pragma.
+      if (auto *FE = SM.getFileEntryForID(SM.getFileID(Range.getBegin())))
+        Out->IWYUPublic.insert(
+            {FE->getLastRef().getUniqueID(),
+             Pragma->startswith("<") || Pragma->startswith("\"")
+                 ? Pragma->str()
+                 : ("\"" + *Pragma + "\"").str()});
+      return false;
+    }
+
+    if (InMainFile) {
+      if (!Pragma->startswith("keep"))
+        return false;
+      // 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 next InclusionDirective is
+      // called, it will know that the next inclusion is behind the IWYU pragma.
+      LastPragmaKeepInMainFileLine = SM.getLineNumber(
+          SM.getMainFileID(), SM.getFileOffset(Range.getBegin()));
+    }
+    return false;
+  }
+
+private:
+  bool InMainFile = false;
+  const SourceManager &SM;
+  PragmaIncludes *Out;
+  // Track the last line "IWYU pragma: keep" was seen in the main file, 1-based.
+  int LastPragmaKeepInMainFileLine = -1;
+};
+
+void PragmaIncludes::record(const CompilerInstance &CI) {
+  auto Record = std::make_unique<RecordPragma>(CI, this);
+  CI.getPreprocessor().addCommentHandler(Record.get());
+  CI.getPreprocessor().addPPCallbacks(std::move(Record));
+}
+
+llvm::StringRef PragmaIncludes::getPublic(const FileEntry *F) const {
+  auto It = IWYUPublic.find(F->getUniqueID());
+  if (It == IWYUPublic.end())
+    return "";
+  return It->getSecond();
+}
+
 std::unique_ptr<ASTConsumer> RecordedAST::record() {
   class Recorder : public ASTConsumer {
     RecordedAST *Out;
 
   public:
     Recorder(RecordedAST *Out) : Out(Out) {}
     void Initialize(ASTContext &Ctx) override { Out->Ctx = &Ctx; }
     bool HandleTopLevelDecl(DeclGroupRef DG) override {

clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

 //===-- RecordTest.cpp ----------------------------------------------------===//
 //
 // 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 "clang-include-cleaner/Record.h"
 #include "clang/Frontend/FrontendAction.h"
+#include "clang/Frontend/FrontendActions.h"
 #include "clang/Testing/TestAST.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang::include_cleaner {
 namespace {
 
TEST_F(RecordASTTest, Macros) {
   Inputs.Code = R"cpp(
     #include "header.h"
     X
   )cpp";
   auto AST = build();
   EXPECT_THAT(Recorded.Roots, testing::ElementsAre(Named("x")));
 }
 
+class PragmaIncludeTest : public ::testing::Test {
+protected:
+  // We don't build an AST, we just run a preprocessor action!
+  TestInputs Inputs;
+  PragmaIncludes PI;
+
+  PragmaIncludeTest() {
+    Inputs.MakeAction = [this] {
+      struct Hook : public PreprocessOnlyAction {
+      public:
+        Hook(PragmaIncludes *Out) : Out(Out) {}
+        bool BeginSourceFileAction(clang::CompilerInstance &CI) override {
+          Out->record(CI);
+          return true;
+        }
+        PragmaIncludes *Out;
+      };
+      return std::make_unique<Hook>(&PI);
+    };
+  }
+
+  TestAST build() { return TestAST(Inputs); }
+};
+
+TEST_F(PragmaIncludeTest, IWYUKeep) {
+  Inputs.Code = R"cpp(// Line 1
+    #include "keep1.h" // IWYU pragma: keep
+    #include "keep2.h" /* IWYU pragma: keep */
+    #include "normal.h"
+  )cpp";
+  Inputs.ExtraFiles["keep1.h"] = Inputs.ExtraFiles["keep2.h"] =
+      Inputs.ExtraFiles["normal.h"] = "";
+
+  TestAST Processed = build();
+  EXPECT_FALSE(PI.shouldKeep(1));
+  EXPECT_TRUE(PI.shouldKeep(2));
+  EXPECT_TRUE(PI.shouldKeep(3));
+  EXPECT_FALSE(PI.shouldKeep(4));
+}
+
+TEST_F(PragmaIncludeTest, IWYUPrivate) {
+  Inputs.Code = R"cpp(
+    #include "public.h"
+  )cpp";
+  Inputs.ExtraFiles["public.h"] = "#include \"private.h\"";
+  Inputs.ExtraFiles["private.h"] = R"cpp(
+    // IWYU pragma: private, include "public2.h"
+    class Private {};
+  )cpp";
+  TestAST Processed = build();
+  auto PrivateFE = Processed.fileManager().getFile("private.h");
+  assert(PrivateFE);
+  EXPECT_EQ(PI.getPublic(PrivateFE.get()), "\"public2.h\"");
+  auto PublicFE = Processed.fileManager().getFile("public.h");
+  assert(PublicFE);
+  EXPECT_EQ(PI.getPublic(PublicFE.get()), ""); // no mapping.
+}
+
 } // namespace
 } // namespace clang::include_cleaner