[include-cleaner] Handle files with unnamed buffers

Some tools can register virtual buffers without identifiers into the filemanager. Make sure we can handle pragmas in such cases. Differential Revision: https://reviews.llvm.org/D157078
llvm · Aug 4, 2023 · 31873d3 · 31873d3
1 parent 734da23
commit 31873d3
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 16 deletions.
diff --git a/clang-tools-extra/include-cleaner/lib/Record.cpp b/clang-tools-extra/include-cleaner/lib/Record.cpp
@@ -279,20 +279,6 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
 
     auto [CommentFID, CommentOffset] = SM.getDecomposedLoc(Range.getBegin());
     int CommentLine = SM.getLineNumber(CommentFID, CommentOffset);
-    auto Filename = SM.getBufferName(Range.getBegin());
-    // Record export pragma.
-    if (Pragma->startswith("export")) {
-      ExportStack.push_back({CommentLine, CommentFID, save(Filename), false});
-    } else if (Pragma->startswith("begin_exports")) {
-      ExportStack.push_back({CommentLine, CommentFID, save(Filename), true});
-    } else if (Pragma->startswith("end_exports")) {
-      // FIXME: be robust on unmatching cases. We should only pop the stack if
-      // the begin_exports and end_exports is in the same file.
-      if (!ExportStack.empty()) {
-        assert(ExportStack.back().Block);
-        ExportStack.pop_back();
-      }
-    }
 
     if (InMainFile) {
       if (Pragma->startswith("keep")) {
@@ -307,8 +293,10 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
 
     auto FE = SM.getFileEntryRefForID(CommentFID);
     if (!FE) {
-      // FIXME: Support IWYU pragmas in virtual files. Our mappings rely on
-      // "persistent" UniqueIDs and that is not the case for virtual files.
+      // This can only happen when the buffer was registered virtually into
+      // SourceManager and FileManager has no idea about it. In such a scenario,
+      // that file cannot be discovered by HeaderSearch, therefore no "explicit"
+      // includes for that file.
       return false;
     }
     auto CommentUID = FE->getUniqueID();
@@ -327,6 +315,20 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
       Out->ShouldKeep.insert(CommentUID);
       return false;
     }
+    auto Filename = FE->getName();
+    // Record export pragma.
+    if (Pragma->startswith("export")) {
+      ExportStack.push_back({CommentLine, CommentFID, save(Filename), false});
+    } else if (Pragma->startswith("begin_exports")) {
+      ExportStack.push_back({CommentLine, CommentFID, save(Filename), true});
+    } else if (Pragma->startswith("end_exports")) {
+      // FIXME: be robust on unmatching cases. We should only pop the stack if
+      // the begin_exports and end_exports is in the same file.
+      if (!ExportStack.empty()) {
+        assert(ExportStack.back().Block);
+        ExportStack.pop_back();
+      }
+    }
     return false;
   }
 

diff --git a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
@@ -7,18 +7,32 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Record.h"
+#include "clang-include-cleaner/Types.h"
+#include "clang/AST/Decl.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "clang/Frontend/FrontendOptions.h"
+#include "clang/Serialization/PCHContainerOperations.h"
 #include "clang/Testing/TestAST.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Annotations/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <cassert>
+#include <memory>
+#include <optional>
+#include <utility>
 
 namespace clang::include_cleaner {
 namespace {
@@ -509,5 +523,38 @@ TEST_F(PragmaIncludeTest, AlwaysKeep) {
   EXPECT_TRUE(PI.shouldKeep(FM.getFile("always_keep.h").get()));
   EXPECT_FALSE(PI.shouldKeep(FM.getFile("usual.h").get()));
 }
+
+TEST_F(PragmaIncludeTest, ExportInUnnamedBuffer) {
+  llvm::StringLiteral Filename = "test.cpp";
+  auto Code = R"cpp(#include "exporter.h")cpp";
+  Inputs.ExtraFiles["exporter.h"] = R"cpp(
+  #pragma once
+  #include "foo.h" // IWYU pragma: export
+  )cpp";
+  Inputs.ExtraFiles["foo.h"] = "";
+
+  auto Clang = std::make_unique<CompilerInstance>(
+      std::make_shared<PCHContainerOperations>());
+  Clang->createDiagnostics();
+
+  Clang->setInvocation(std::make_unique<CompilerInvocation>());
+  ASSERT_TRUE(CompilerInvocation::CreateFromArgs(
+      Clang->getInvocation(), {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("foo.h")), *FM),
+      testing::ElementsAre(llvm::cantFail(FM->getFileRef("exporter.h"))));
+}
 } // namespace
 } // namespace clang::include_cleaner