[clangd] IncludeCleaner: Handle macros coming from ScratchBuffer

Reviewed By: sammccall

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

clang-tools-extra/clangd/IncludeCleaner.cpp

@@ -121,8 +121,14 @@ struct ReferencedFiles {
     if (!Macros.insert(FID).second)
       return;
     const auto &Exp = SM.getSLocEntry(FID).getExpansion();
+    // For token pasting operator in macros, spelling and expansion locations
+    // can be within a temporary buffer that Clang creates (scratch space or
+    // ScratchBuffer). That is not a real file we can include.
+    if (!SM.isWrittenInScratchSpace(Exp.getSpellingLoc()))
       add(Exp.getSpellingLoc());
+    if (!SM.isWrittenInScratchSpace(Exp.getExpansionLocStart()))
       add(Exp.getExpansionLocStart());
+    if (!SM.isWrittenInScratchSpace(Exp.getExpansionLocEnd()))
       add(Exp.getExpansionLocEnd());
   }
 };

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

@@ -167,6 +167,47 @@ TEST(IncludeCleaner, GetUnusedHeaders) {
               UnorderedElementsAre("\"unused.h\"", "\"dir/unused.h\""));
 }
 
+TEST(IncludeCleaner, ScratchBuffer) {
+  TestTU TU;
+  TU.Filename = "foo.cpp";
+  TU.Code = R"cpp(
+    #include "macro_spelling_in_scratch_buffer.h"
+
+    using flags::FLAGS_FOO;
+
+    int concat(a, b) = 42;
+    )cpp";
+  // The pasting operator in combination with DEFINE_FLAG will create
+  // ScratchBuffer with `flags::FLAGS_FOO` that will have FileID but not
+  // FileEntry.
+  TU.AdditionalFiles["macro_spelling_in_scratch_buffer.h"] = R"cpp(
+    #define DEFINE_FLAG(X) \
+    namespace flags { \
+    int FLAGS_##X; \
+    } \
+
+    DEFINE_FLAG(FOO)
+
+    #define ab x
+    #define concat(x, y) x##y
+    )cpp";
+  ParsedAST AST = TU.build();
+  auto &SM = AST.getSourceManager();
+  auto &Includes = AST.getIncludeStructure();
+  auto ReferencedFiles = findReferencedFiles(findReferencedLocations(AST), SM);
+  auto Entry = SM.getFileManager().getFile(
+      testPath("macro_spelling_in_scratch_buffer.h"));
+  ASSERT_TRUE(Entry);
+  auto FID = SM.translateFile(*Entry);
+  // No "<scratch space>" FID.
+  EXPECT_THAT(ReferencedFiles, UnorderedElementsAre(FID));
+  // Should not crash due to <scratch space> "files" missing from include
+  // structure.
+  EXPECT_THAT(
+      getUnused(Includes, translateToHeaderIDs(ReferencedFiles, Includes, SM)),
+      ::testing::IsEmpty());
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang