[include-cleaner] Initial version for the "Location=>Header" step

This patch implements the initial version of "Location => Header" step:

- define the interface;
- integrate into the existing workflow, and use the PragmaIncludes;

Differential Revision: https://reviews.llvm.org/D137320

7 files changed, 163 insertions, 35 deletions

diff --git a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
index 62277199871b..7b8606f4a6ce 100644
--- a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
+++ b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
@@ -11,6 +11,7 @@
 #ifndef CLANG_INCLUDE_CLEANER_ANALYSIS_H
 #define CLANG_INCLUDE_CLEANER_ANALYSIS_H
 
+#include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
@@ -42,11 +43,9 @@ using UsedSymbolCB = llvm::function_ref<void(const SymbolReference &SymRef,
 ///    headers which don't match any #include in the main file
 ///  - to diagnose unused includes: an #include in the main file does not match
 ///    the headers for any referenced symbol
-/// FIXME: Take in an include structure to improve location to header mappings
-/// (e.g. IWYU pragmas).
 void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
-              llvm::ArrayRef<SymbolReference> MacroRefs, const SourceManager &,
-              UsedSymbolCB CB);
+              llvm::ArrayRef<SymbolReference> MacroRefs,
+              const PragmaIncludes &PI, const SourceManager &, UsedSymbolCB CB);
 
 } // namespace include_cleaner
 } // namespace clang
diff --git a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h
index 6aaf6cfa5d37..87db0cf585a3 100644
--- a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h
+++ b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h
@@ -105,10 +105,14 @@ struct Header {
     Physical,
     /// A recognized standard library header, like <string>.
     Standard,
+    /// A verbatim header spelling, a string quoted with <> or "" that can be
+    /// #included directly.
+    VerbatimSpelling,
   };
 
   Header(const FileEntry *FE) : Storage(FE) {}
   Header(tooling::stdlib::Header H) : Storage(H) {}
+  Header(StringRef VerbatimSpelling) : Storage(VerbatimSpelling) {}
 
   Kind kind() const { return static_cast<Kind>(Storage.index()); }
   bool operator==(const Header &RHS) const { return Storage == RHS.Storage; }
@@ -117,11 +121,13 @@ struct Header {
   tooling::stdlib::Header standard() const {
     return std::get<Standard>(Storage);
   }
+  StringRef verbatimSpelling() const {
+    return std::get<VerbatimSpelling>(Storage);
+  }
 
 private:
-  // FIXME: Handle verbatim spellings.
   // Order must match Kind enum!
-  std::variant<const FileEntry *, tooling::stdlib::Header> Storage;
+  std::variant<const FileEntry *, tooling::stdlib::Header, StringRef> Storage;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Header &);
 
diff --git a/clang-tools-extra/include-cleaner/lib/Analysis.cpp b/clang-tools-extra/include-cleaner/lib/Analysis.cpp
index f7c8bdeee0bf..4d66c34ce788 100644
--- a/clang-tools-extra/include-cleaner/lib/Analysis.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Analysis.cpp
@@ -17,20 +17,11 @@
 #include "llvm/ADT/SmallVector.h"
 
 namespace clang::include_cleaner {
-namespace {
-llvm::SmallVector<Header>
-toHeader(llvm::ArrayRef<tooling::stdlib::Header> Headers) {
-  llvm::SmallVector<Header> Result;
-  llvm::for_each(Headers, [&](tooling::stdlib::Header H) {
-    Result.emplace_back(Header(H));
-  });
-  return Result;
-}
-} // namespace
 
 void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
               llvm::ArrayRef<SymbolReference> MacroRefs,
-              const SourceManager &SM, UsedSymbolCB CB) {
+              const PragmaIncludes &PI, const SourceManager &SM,
+              UsedSymbolCB CB) {
   tooling::stdlib::Recognizer Recognizer;
   for (auto *Root : ASTRoots) {
     auto &SM = Root->getASTContext().getSourceManager();
@@ -38,23 +29,52 @@ void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
       if (auto SS = Recognizer(&ND)) {
         // FIXME: Also report forward decls from main-file, so that the caller
         // can decide to insert/ignore a header.
-        return CB({Loc, Symbol(*SS), RT}, toHeader(SS->headers()));
+        return CB({Loc, Symbol(*SS), RT}, findIncludeHeaders({*SS}, SM, PI));
       }
       // FIXME: Extract locations from redecls.
-      // FIXME: Handle IWYU pragmas, non self-contained files.
-      // FIXME: Handle macro locations.
-      if (auto *FE = SM.getFileEntryForID(SM.getFileID(ND.getLocation())))
-        return CB({Loc, Symbol(ND), RT}, {Header(FE)});
+      return CB({Loc, Symbol(ND), RT},
+                findIncludeHeaders({ND.getLocation()}, SM, PI));
     });
   }
   for (const SymbolReference &MacroRef : MacroRefs) {
     assert(MacroRef.Target.kind() == Symbol::Macro);
-    // FIXME: Handle IWYU pragmas, non self-contained files.
     // FIXME: Handle macro locations.
-    if (auto *FE = SM.getFileEntryForID(
-            SM.getFileID(MacroRef.Target.macro().Definition)))
-      CB(MacroRef, {Header(FE)});
+    return CB(MacroRef,
+              findIncludeHeaders(MacroRef.Target.macro().Definition, SM, PI));
+  }
+}
+
+llvm::SmallVector<Header> findIncludeHeaders(const SymbolLocation &SLoc,
+                                             const SourceManager &SM,
+                                             const PragmaIncludes &PI) {
+  llvm::SmallVector<Header> Results;
+  if (auto *Loc = std::get_if<SourceLocation>(&SLoc)) {
+    // FIXME: Handle non self-contained files.
+    FileID FID = SM.getFileID(*Loc);
+    const auto *FE = SM.getFileEntryForID(FID);
+    if (!FE)
+      return {};
+
+    // We treat the spelling header in the IWYU pragma as the final public
+    // header.
+    // FIXME: look for exporters if the public header is exported by another
+    // header.
+    llvm::StringRef VerbatimSpelling = PI.getPublic(FE);
+    if (!VerbatimSpelling.empty())
+      return {{VerbatimSpelling}};
+
+    Results = {{FE}};
+    // FIXME: compute transitive exporter headers.
+    for (const auto *Export : PI.getExporters(FE, SM.getFileManager()))
+      Results.push_back(Export);
+    return Results;
+  }
+  if (auto *Sym = std::get_if<tooling::stdlib::Symbol>(&SLoc)) {
+    for (const auto &H : Sym->headers())
+      Results.push_back(H);
+    return Results;
   }
+  llvm_unreachable("unhandled SymbolLocation kind!");
 }
 
 } // namespace clang::include_cleaner
diff --git a/clang-tools-extra/include-cleaner/lib/AnalysisInternal.h b/clang-tools-extra/include-cleaner/lib/AnalysisInternal.h
index ff5a49c99410..18bcedde1a2e 100644
--- a/clang-tools-extra/include-cleaner/lib/AnalysisInternal.h
+++ b/clang-tools-extra/include-cleaner/lib/AnalysisInternal.h
@@ -21,6 +21,7 @@
 #ifndef CLANG_INCLUDE_CLEANER_ANALYSISINTERNAL_H
 #define CLANG_INCLUDE_CLEANER_ANALYSISINTERNAL_H
 
+#include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
@@ -46,6 +47,18 @@ namespace include_cleaner {
 void walkAST(Decl &Root,
              llvm::function_ref<void(SourceLocation, NamedDecl &, RefType)>);
 
+/// A location where a symbol can be provided.
+/// It is either a physical file of the TU (SourceLocation) or a logical
+/// location in the standard library (stdlib::Symbol).
+// FIXME: use a real Class!
+using SymbolLocation = std::variant<SourceLocation, tooling::stdlib::Symbol>;
+
+/// Finds the headers that provide the symbol location.
+// FIXME: expose signals
+llvm::SmallVector<Header> findIncludeHeaders(const SymbolLocation &Loc,
+                                             const SourceManager &SM,
+                                             const PragmaIncludes &PI);
+
 /// Write an HTML summary of the analysis to the given stream.
 /// FIXME: Once analysis has a public API, this should be public too.
 void writeHTMLReport(FileID File, llvm::ArrayRef<Decl *> Roots, ASTContext &Ctx,
diff --git a/clang-tools-extra/include-cleaner/lib/Record.cpp b/clang-tools-extra/include-cleaner/lib/Record.cpp
index 948b14c3338f..d824e5d06bad 100644
--- a/clang-tools-extra/include-cleaner/lib/Record.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Record.cpp
@@ -332,6 +332,10 @@ RecordedPP::RecordedIncludes::match(Header H) const {
     for (unsigned I : BySpelling.lookup(H.standard().name().trim("<>")))
       Result.push_back(&All[I]);
     break;
+  case Header::VerbatimSpelling:
+    for (unsigned I : BySpelling.lookup(H.verbatimSpelling().trim("\"<>")))
+      Result.push_back(&All[I]);
+    break;
   }
   return Result;
 }
diff --git a/clang-tools-extra/include-cleaner/lib/Types.cpp b/clang-tools-extra/include-cleaner/lib/Types.cpp
index ed1509295896..340843e6bdf9 100644
--- a/clang-tools-extra/include-cleaner/lib/Types.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Types.cpp
@@ -34,6 +34,8 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Header &H) {
     return OS << H.physical()->getName();
   case Header::Standard:
     return OS << H.standard().name();
+  case Header::VerbatimSpelling:
+    return OS << H.verbatimSpelling();
   }
   llvm_unreachable("Unhandled Header kind");
 }
diff --git a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
index 49652ad35023..594a0215c95a 100644
--- a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
+++ b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
@@ -7,11 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Analysis.h"
+#include "AnalysisInternal.h"
+#include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Frontend/FrontendActions.h"
 #include "clang/Testing/TestAST.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -28,19 +31,39 @@ using testing::UnorderedElementsAre;
 
 TEST(WalkUsed, Basic) {
   // FIXME: Have a fixture for setting up tests.
-  llvm::Annotations HeaderCode(R"cpp(
-  void foo();
-  namespace std { class vector {}; })cpp");
   llvm::Annotations Code(R"cpp(
-  void $bar^bar() {
+  #include "header.h"
+  #include "private.h"
+
+  void $bar^bar($private^Private) {
     $foo^foo();
     std::$vector^vector $vconstructor^v;
   }
   )cpp");
   TestInputs Inputs(Code.code());
-  Inputs.ExtraFiles["header.h"] = HeaderCode.code().str();
-  Inputs.ExtraArgs.push_back("-include");
-  Inputs.ExtraArgs.push_back("header.h");
+  Inputs.ExtraFiles["header.h"] = R"cpp(
+  void foo();
+  namespace std { class vector {}; }
+  )cpp";
+  Inputs.ExtraFiles["private.h"] = R"cpp(
+    // IWYU pragma: private, include "path/public.h"
+    class Private {};
+  )cpp";
+
+  PragmaIncludes PI;
+  Inputs.MakeAction = [&PI] {
+    struct Hook : public SyntaxOnlyAction {
+    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 AST(Inputs);
 
   llvm::SmallVector<Decl *> TopLevelDecls;
@@ -50,7 +73,7 @@ TEST(WalkUsed, Basic) {
 
   auto &SM = AST.sourceManager();
   llvm::DenseMap<size_t, std::vector<Header>> OffsetToProviders;
-  walkUsed(TopLevelDecls, /*MacroRefs=*/{}, SM,
+  walkUsed(TopLevelDecls, /*MacroRefs=*/{}, PI, SM,
            [&](const SymbolReference &Ref, llvm::ArrayRef<Header> Providers) {
              auto [FID, Offset] = SM.getDecomposedLoc(Ref.RefLocation);
              EXPECT_EQ(FID, SM.getMainFileID());
@@ -63,6 +86,8 @@ TEST(WalkUsed, Basic) {
       OffsetToProviders,
       UnorderedElementsAre(
           Pair(Code.point("bar"), UnorderedElementsAre(MainFile)),
+          Pair(Code.point("private"),
+               UnorderedElementsAre(Header("\"path/public.h\""))),
           Pair(Code.point("foo"), UnorderedElementsAre(HeaderFile)),
           Pair(Code.point("vector"), UnorderedElementsAre(VectorSTL)),
           Pair(Code.point("vconstructor"), UnorderedElementsAre(VectorSTL))));
@@ -89,10 +114,11 @@ TEST(WalkUsed, MacroRefs) {
   Symbol Answer =
       Macro{&Idents.get("ANSWER"), SM.getComposedLoc(HdrID, Hdr.point())};
   llvm::DenseMap<size_t, std::vector<Header>> OffsetToProviders;
+  PragmaIncludes PI;
   walkUsed(/*ASTRoots=*/{}, /*MacroRefs=*/
            {SymbolReference{SM.getComposedLoc(SM.getMainFileID(), Main.point()),
                             Answer, RefType::Explicit}},
-           SM,
+           PI, SM,
            [&](const SymbolReference &Ref, llvm::ArrayRef<Header> Providers) {
              auto [FID, Offset] = SM.getDecomposedLoc(Ref.RefLocation);
              EXPECT_EQ(FID, SM.getMainFileID());
@@ -104,5 +130,63 @@ TEST(WalkUsed, MacroRefs) {
       UnorderedElementsAre(Pair(Main.point(), UnorderedElementsAre(HdrFile))));
 }
 
+TEST(FindIncludeHeaders, IWYU) {
+  TestInputs Inputs;
+  PragmaIncludes PI;
+  Inputs.MakeAction = [&PI] {
+    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);
+  };
+
+  Inputs.Code = R"cpp(
+    #include "header1.h"
+    #include "header2.h"
+  )cpp";
+  Inputs.ExtraFiles["header1.h"] = R"cpp(
+    // IWYU pragma: private, include "path/public.h"
+  )cpp";
+  Inputs.ExtraFiles["header2.h"] = R"cpp(
+    #include "detail1.h" // IWYU pragma: export
+
+    // IWYU pragma: begin_exports
+    #include "detail2.h"
+    // IWYU pragma: end_exports
+
+    #include "normal.h"
+  )cpp";
+  Inputs.ExtraFiles["normal.h"] = Inputs.ExtraFiles["detail1.h"] =
+      Inputs.ExtraFiles["detail2.h"] = "";
+  TestAST AST(Inputs);
+  const auto &SM = AST.sourceManager();
+  auto &FM = SM.getFileManager();
+  // Returns the source location for the start of the file.
+  auto SourceLocFromFile = [&](llvm::StringRef FileName) {
+    return SM.translateFileLineCol(FM.getFile(FileName).get(),
+                                   /*Line=*/1, /*Col=*/1);
+  };
+
+  EXPECT_THAT(findIncludeHeaders(SourceLocFromFile("header1.h"), SM, PI),
+              UnorderedElementsAre(Header("\"path/public.h\"")));
+
+  EXPECT_THAT(findIncludeHeaders(SourceLocFromFile("detail1.h"), SM, PI),
+              UnorderedElementsAre(Header(FM.getFile("header2.h").get()),
+                                   Header(FM.getFile("detail1.h").get())));
+  EXPECT_THAT(findIncludeHeaders(SourceLocFromFile("detail2.h"), SM, PI),
+              UnorderedElementsAre(Header(FM.getFile("header2.h").get()),
+                                   Header(FM.getFile("detail2.h").get())));
+
+  EXPECT_THAT(findIncludeHeaders(SourceLocFromFile("normal.h"), SM, PI),
+              UnorderedElementsAre(Header(FM.getFile("normal.h").get())));
+}
+
 } // namespace
 } // namespace clang::include_cleaner