[include-cleaner] FindHeaders respects IWYU export pragma for standard headers.

[hokein]

Summary

Fixes https://github.com/llvm/llvm-project/issues/59927

[kadircet]

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

206

just saying StandardHeader here isn't really useful. what about changing signautre here to take a include_cleaner::Header instead, named IncludedHeader and switch over kind to add it to ExportBy map? I don't think there's any value in storing Exporters for <vector> both as a physical entry and a stdlib entry.

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

122

nit: tooling::stdlib::Header::named("<string>")

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

437

drop the comment Line 1

[hokein]

address comments

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

206

Ah, nice idea! thanks.

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

122

We need this tooling::stdlib::Symbol here, include_cleaner::findHeaders accepts a SymbolLocation, SymbolLocation is construct from a tooling::stdlib::Symbol.

[kadircet]

thanks!

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

192–200

nit: I'd make this an optional instead (as the code in checkForExport looks weird now, we do nullptr check on one branch but not the other) and change the flow to:

std::optional<Header> IncludedHeader;
if(IsAngled) {
  if(stdlibheader) ...
}

if(!IncludedHeader && File) {
 ...
}

214–227

nit:

switch(IncludedHeader.kind()) {
case Verbatim:
  assert(false && "...");
  break;
case others:
...
}

to be consistent with rest of the places (and more explicit about handling of verbatim)

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

122

oops you're right

[hokein]

address comments.

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

public:
   /// Returns the public mapping include for the given physical header file.
   /// Returns "" if there is none.
   llvm::StringRef getPublic(const FileEntry *File) const;
 
   /// Returns all direct exporter headers for the given header file.
   /// 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;
 
   /// Returns true if the given file is marked with the IWYU private pragma.
   bool isPrivate(const FileEntry *File) const;
 
 private:
private:
   /// hasn't been opened before. So store the path and convert it to a
   /// FileEntry by opening the file again through a FileManager.
   ///
   /// We don't use RealPathName, as opening the file through a different name
   /// changes its preferred name. Clearly this is fragile!
   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;
 
   /// Owns the strings.
   llvm::BumpPtrAllocator Arena;
 
   // FIXME: add support for clang use_instead pragma

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

while (FE) {
 
       // Walkup the include stack for non self-contained headers.
       FID = SM.getDecomposedIncludedLoc(FID).first;
       FE = SM.getFileEntryForID(FID);
     }
     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;
   }
   }
   llvm_unreachable("unhandled SymbolLocation kind!");
 }
 
 } // namespace clang::include_cleaner

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

 #include "clang/AST/DeclGroup.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/MacroInfo.h"
 #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 {
 
 class PPRecorder : public PPCallbacks {
 public:
   PPRecorder(RecordedPP &Recorded, const Preprocessor &PP)
       : Recorded(Recorded), PP(PP), SM(PP.getSourceManager()) {}
void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                           CharSourceRange /*FilenameRange*/,
                           OptionalFileEntryRef File,
                           llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const clang::Module * /*Imported*/,
                           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();
     if (Top.SeenAtFile != IncludingFile)
       return;
     // 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);
     }
     if (!Top.Block) // Pop immediately for single-line export pragma.
       ExportStack.pop_back();
   }
 

 
 llvm::StringRef PragmaIncludes::getPublic(const FileEntry *F) const {
   auto It = IWYUPublic.find(F->getUniqueID());
   if (It == IWYUPublic.end())
     return "";
   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 {
   return !NonSelfContainedFiles.contains(FE->getUniqueID());
 }
 
 bool PragmaIncludes::isPrivate(const FileEntry *FE) const {
   return IWYUPublic.find(FE->getUniqueID()) != IWYUPublic.end();

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

 #include "clang-include-cleaner/Analysis.h"
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/FileEntry.h"
 #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"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <memory>
 
 namespace clang::include_cleaner {
EXPECT_THAT(findHeaders("exported2.h"),
               UnorderedElementsAre(physicalHeader("exported2.h"),
                                    physicalHeader("exporter.h")));
   EXPECT_THAT(findHeaders("normal.h"),
               UnorderedElementsAre(physicalHeader("normal.h")));
   EXPECT_THAT(findHeaders("exporter.h"),
               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"
   )cpp";
   Inputs.ExtraFiles["header.h"] = guard(R"cpp(
     #include "fragment.inc"
   )cpp");
   Inputs.ExtraFiles["fragment.inc"] = "";

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/Basic/SourceLocation.h"
 #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"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang::include_cleaner {
 namespace {
TEST_F(PragmaIncludeTest, IWYUExport) {
 
   EXPECT_TRUE(PI.getExporters(FM.getFile("export1.h").get(), FM).empty());
   EXPECT_TRUE(PI.getExporters(FM.getFile("export2.h").get(), FM).empty());
   EXPECT_TRUE(PI.getExporters(FM.getFile("export3.h").get(), FM).empty());
   EXPECT_TRUE(
       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"
   )cpp";
   Inputs.ExtraFiles["normal.h"] = R"cpp(
     #include "foo.h"
 
     // IWYU pragma: begin_exports