[include-cleaner] Treat include_next as exporting

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

@@ -17,6 +17,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
+#include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/DirectoryLookup.h"
 #include "clang/Lex/MacroInfo.h"
@@ -27,19 +28,15 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Allocator.h"
-#include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
 #include "llvm/Support/StringSaver.h"
 #include <algorithm>
 #include <assert.h>
 #include <memory>
 #include <optional>
-#include <set>
 #include <utility>
 #include <vector>
 
@@ -225,35 +222,41 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
       }
     if (!IncludedHeader && File)
       IncludedHeader = *File;
-    checkForExport(HashFID, HashLine, std::move(IncludedHeader), File);
+    checkForExport(HashFID, HashLine, std::move(IncludedHeader), File,
+                   IncludeTok);
     checkForKeep(HashLine, File);
   }
 
   void checkForExport(FileID IncludingFile, int HashLine,
                       std::optional<Header> IncludedHeader,
-                      OptionalFileEntryRef IncludedFile) {
-    if (ExportStack.empty())
+                      OptionalFileEntryRef IncludedFile,
+                      const Token &IncludeTok) {
+    // Adds an export for the IncludedFile from IncludingFile.
+    auto AddExport = [&]() {
+      auto ExportingFileName = SM.getFileEntryForID(IncludingFile)->getName();
+      if (IncludedFile) {
+        Out->IWYUExportBy[IncludedFile->getFileEntry().getUniqueID()].push_back(
+            ExportingFileName);
+      }
+      if (IncludedHeader && IncludedHeader->kind() == Header::Standard) {
+        Out->StdIWYUExportBy[IncludedHeader->standard()].push_back(
+            ExportingFileName);
+      }
+    };
+    if (ExportStack.empty()) {
+      if (IncludeTok.getIdentifierInfo()->getPPKeywordID() ==
+          tok::pp_include_next) {
+        AddExport();
+      }
       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) {
-        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;
-        }
-      }
+      AddExport();
       // main-file #include with export pragma should never be removed.
       if (Top.SeenAtFile == SM.getMainFileID() && IncludedFile)
         Out->ShouldKeep.insert(IncludedFile->getUniqueID());

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

@@ -452,6 +452,8 @@ TEST_F(PragmaIncludeTest, IWYUExportForStandardHeaders) {
   auto &FM = Processed.fileManager();
   EXPECT_THAT(PI.getExporters(*tooling::stdlib::Header::named("<string>"), FM),
               testing::UnorderedElementsAre(FileNamed("export.h")));
+  EXPECT_THAT(PI.getExporters(llvm::cantFail(FM.getFileRef("string")), FM),
+              testing::UnorderedElementsAre(FileNamed("export.h")));
 }
 
 TEST_F(PragmaIncludeTest, IWYUExportBlock) {
@@ -556,5 +558,37 @@ TEST_F(PragmaIncludeTest, ExportInUnnamedBuffer) {
       PI.getExporters(llvm::cantFail(FM->getFileRef("foo.h")), *FM),
       testing::ElementsAre(llvm::cantFail(FM->getFileRef("exporter.h"))));
 }
+
+TEST_F(PragmaIncludeTest, ExportIncludeNext) {
+  llvm::StringLiteral Filename = "test.cpp";
+  auto Code = R"cpp(#include <new>)cpp";
+  Inputs.ExtraFiles["foo/new"] = R"cpp(
+  #pragma once
+  #include_next <new>
+  )cpp";
+  Inputs.ExtraFiles["inner/new"] = "#pragma once";
+
+  auto Clang = std::make_unique<CompilerInstance>(
+      std::make_shared<PCHContainerOperations>());
+  Clang->createDiagnostics();
+
+  Clang->setInvocation(std::make_unique<CompilerInvocation>());
+  ASSERT_TRUE(CompilerInvocation::CreateFromArgs(
+      Clang->getInvocation(), {"-Ifoo/", "-Iinner/", Filename.data()},
+      Clang->getDiagnostics(), "clang"));
+
+  // Create unnamed memory buffers for all the files.
+  auto VFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  VFS->addFile(Filename, /*ModificationTime=*/0,
+               llvm::MemoryBuffer::getMemBufferCopy(Code, /*BufferName=*/""));
+  for (const auto &Extra : Inputs.ExtraFiles)
+    VFS->addFile(Extra.getKey(), /*ModificationTime=*/0,
+                 llvm::MemoryBuffer::getMemBufferCopy(Extra.getValue(),
+                                                      /*BufferName=*/""));
+  auto *FM = Clang->createFileManager(VFS);
+  ASSERT_TRUE(Clang->ExecuteAction(*Inputs.MakeAction()));
+  EXPECT_THAT(PI.getExporters(llvm::cantFail(FM->getFileRef("inner/new")), *FM),
+              testing::ElementsAre(llvm::cantFail(FM->getFileRef("foo/new"))));
+}
 } // namespace
 } // namespace clang::include_cleaner