[include-cleaner] FindHeaders respects IWYU export pragma for standar…

…d headers.

Fixes #59927

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

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

@@ -66,6 +66,8 @@ class PragmaIncludes {
   /// Returns empty if there is none.
   llvm::SmallVector<const FileEntry *> getExporters(const FileEntry *File,
                                                     FileManager &FM) const;
+  llvm::SmallVector<const FileEntry *> getExporters(tooling::stdlib::Header,
+                                                    FileManager &FM) const;
 
   /// Returns true if the given file is a self-contained file.
   bool isSelfContained(const FileEntry *File) const;
@@ -100,6 +102,9 @@ class PragmaIncludes {
   llvm::DenseMap<llvm::sys::fs::UniqueID,
                  llvm::SmallVector</*FileEntry::getName()*/ llvm::StringRef>>
       IWYUExportBy;
+  llvm::DenseMap<tooling::stdlib::Header,
+                 llvm::SmallVector</*FileEntry::getName()*/ llvm::StringRef>>
+      StdIWYUExportBy;
 
   /// Contains all non self-contained files detected during the parsing.
   llvm::DenseSet<llvm::sys::fs::UniqueID> NonSelfContainedFiles;

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

@@ -45,8 +45,11 @@ llvm::SmallVector<Header> findHeaders(const SymbolLocation &Loc,
     return Results;
   }
   case SymbolLocation::Standard: {
-    for (const auto &H : Loc.standard().headers())
+    for (const auto &H : Loc.standard().headers()) {
       Results.push_back(H);
+      for (const auto *Export : PI->getExporters(H, SM.getFileManager()))
+        Results.push_back(Header(Export));
+    }
     return Results;
   }
   }

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

@@ -18,6 +18,7 @@
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Inclusions/HeaderAnalysis.h"
+#include "clang/Tooling/Inclusions/StandardLibrary.h"
 
 namespace clang::include_cleaner {
 namespace {
@@ -188,12 +189,20 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
                           SrcMgr::CharacteristicKind FileKind) override {
     FileID HashFID = SM.getFileID(HashLoc);
     int HashLine = SM.getLineNumber(HashFID, SM.getFileOffset(HashLoc));
-    checkForExport(HashFID, HashLine, File ? &File->getFileEntry() : nullptr);
+    std::optional<Header> IncludedHeader;
+    if (IsAngled)
+      if (auto StandardHeader =
+              tooling::stdlib::Header::named("<" + FileName.str() + ">")) {
+        IncludedHeader = *StandardHeader;
+      }
+    if (!IncludedHeader && File)
+      IncludedHeader = &File->getFileEntry();
+    checkForExport(HashFID, HashLine, std::move(IncludedHeader));
     checkForKeep(HashLine);
   }
 
   void checkForExport(FileID IncludingFile, int HashLine,
-                      const FileEntry *IncludedHeader) {
+                      std::optional<Header> IncludedHeader) {
     if (ExportStack.empty())
       return;
     auto &Top = ExportStack.back();
@@ -202,9 +211,20 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
     // Make sure current include is covered by the export pragma.
     if ((Top.Block && HashLine > Top.SeenAtLine) ||
         Top.SeenAtLine == HashLine) {
-      if (IncludedHeader)
-        Out->IWYUExportBy[IncludedHeader->getUniqueID()].push_back(
-            Top.Path);
+      if (IncludedHeader) {
+        switch (IncludedHeader->kind()) {
+        case Header::Physical:
+          Out->IWYUExportBy[IncludedHeader->physical()->getUniqueID()]
+              .push_back(Top.Path);
+          break;
+        case Header::Standard:
+          Out->StdIWYUExportBy[IncludedHeader->standard()].push_back(Top.Path);
+          break;
+        case Header::Verbatim:
+          assert(false && "unexpected Verbatim header");
+          break;
+        }
+      }
       // main-file #include with export pragma should never be removed.
       if (Top.SeenAtFile == SM.getMainFileID())
         Out->ShouldKeep.insert(HashLine);
@@ -329,19 +349,32 @@ llvm::StringRef PragmaIncludes::getPublic(const FileEntry *F) const {
   return It->getSecond();
 }
 
+static llvm::SmallVector<const FileEntry *>
+toFileEntries(llvm::ArrayRef<StringRef> FileNames, FileManager& FM) {
+    llvm::SmallVector<const FileEntry *> Results;
+
+  for (auto FName : FileNames) {
+    // FIMXE: log the failing cases?
+    if (auto FE = expectedToOptional(FM.getFileRef(FName)))
+      Results.push_back(*FE);
+  }
+  return Results;
+}
 llvm::SmallVector<const FileEntry *>
 PragmaIncludes::getExporters(const FileEntry *File, FileManager &FM) const {
   auto It = IWYUExportBy.find(File->getUniqueID());
   if (It == IWYUExportBy.end())
     return {};
 
-  llvm::SmallVector<const FileEntry *> Results;
-  for (auto Export : It->getSecond()) {
-    // FIMXE: log the failing cases?
-    if (auto FE = expectedToOptional(FM.getFileRef(Export)))
-      Results.push_back(*FE);
-  }
-  return Results;
+  return toFileEntries(It->getSecond(), FM);
+}
+llvm::SmallVector<const FileEntry *>
+PragmaIncludes::getExporters(tooling::stdlib::Header StdHeader,
+                             FileManager &FM) const {
+  auto It = StdIWYUExportBy.find(StdHeader);
+  if (It == StdIWYUExportBy.end())
+    return {};
+  return toFileEntries(It->getSecond(), FM);
 }
 
 bool PragmaIncludes::isSelfContained(const FileEntry *FE) const {

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

@@ -15,6 +15,7 @@
 #include "clang/Basic/FileManager.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Testing/TestAST.h"
+#include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Annotations/Annotations.h"
@@ -107,6 +108,23 @@ TEST_F(FindHeadersTest, IWYUExport) {
               UnorderedElementsAre(physicalHeader("exporter.h")));
 }
 
+TEST_F(FindHeadersTest, IWYUExportForStandardHeaders) {
+  Inputs.Code = R"cpp(
+    #include "exporter.h"
+  )cpp";
+  Inputs.ExtraFiles["exporter.h"] = guard(R"cpp(
+    #include <string> // IWYU pragma: export
+  )cpp");
+  Inputs.ExtraFiles["string"] = guard("");
+  Inputs.ExtraArgs.push_back("-isystem.");
+  buildAST();
+  tooling::stdlib::Symbol StdString =
+      *tooling::stdlib::Symbol::named("std::", "string");
+  EXPECT_THAT(
+      include_cleaner::findHeaders(StdString, AST->sourceManager(), &PI),
+      UnorderedElementsAre(physicalHeader("exporter.h"), StdString.header()));
+}
+
 TEST_F(FindHeadersTest, SelfContained) {
   Inputs.Code = R"cpp(
     #include "header.h"

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

@@ -11,6 +11,7 @@
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Testing/TestAST.h"
+#include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Annotations/Annotations.h"
@@ -432,6 +433,21 @@ TEST_F(PragmaIncludeTest, IWYUExport) {
       PI.getExporters(SM.getFileEntryForID(SM.getMainFileID()), FM).empty());
 }
 
+TEST_F(PragmaIncludeTest, IWYUExportForStandardHeaders) {
+  Inputs.Code = R"cpp(
+    #include "export.h"
+  )cpp";
+  Inputs.ExtraFiles["export.h"] = R"cpp(
+    #include <string> // IWYU pragma: export
+  )cpp";
+  Inputs.ExtraFiles["string"] = "";
+  Inputs.ExtraArgs = {"-isystem."};
+  TestAST Processed = build();
+  auto &FM = Processed.fileManager();
+  EXPECT_THAT(PI.getExporters(*tooling::stdlib::Header::named("<string>"), FM),
+              testing::UnorderedElementsAre(FileNamed("export.h")));
+}
+
 TEST_F(PragmaIncludeTest, IWYUExportBlock) {
   Inputs.Code = R"cpp(// Line 1
    #include "normal.h"