VitaNuo updated this revision to Diff 532964.
VitaNuo marked 8 inline comments as done.
VitaNuo added a comment.
Herald added a subscriber: javed.absar.

Address review comments.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D152900/new/

https://reviews.llvm.org/D152900

Files:
  clang-tools-extra/clangd/ClangdServer.cpp
  clang-tools-extra/clangd/ParsedAST.cpp
  clang-tools-extra/clangd/Preamble.cpp
  clang-tools-extra/clangd/Preamble.h
  clang-tools-extra/clangd/TUScheduler.cpp
  clang-tools-extra/clangd/TUScheduler.h
  clang-tools-extra/clangd/index/CanonicalIncludes.cpp
  clang-tools-extra/clangd/index/CanonicalIncludes.h
  clang-tools-extra/clangd/index/FileIndex.cpp
  clang-tools-extra/clangd/index/FileIndex.h
  clang-tools-extra/clangd/index/IndexAction.cpp
  clang-tools-extra/clangd/index/SymbolCollector.cpp
  clang-tools-extra/clangd/index/SymbolCollector.h
  clang-tools-extra/clangd/tool/Check.cpp
  clang-tools-extra/clangd/unittests/CanonicalIncludesTests.cpp
  clang-tools-extra/clangd/unittests/FileIndexTests.cpp
  clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
  clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
  clang-tools-extra/clangd/unittests/TestTU.cpp
  clang-tools-extra/clangd/unittests/TestWorkspace.cpp

Index: clang-tools-extra/clangd/unittests/TestWorkspace.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/TestWorkspace.cpp
+++ clang-tools-extra/clangd/unittests/TestWorkspace.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "TestWorkspace.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/FileIndex.h"
 #include "gtest/gtest.h"
 #include <optional>
@@ -22,9 +23,10 @@
     TU.Code = Input.second.Code;
     TU.Filename = Input.first().str();
     TU.preamble([&](ASTContext &Ctx, Preprocessor &PP,
-                    const CanonicalIncludes &CanonIncludes) {
+                    const CanonicalIncludes &CanonIncludes,
+                    const include_cleaner::PragmaIncludes *PI) {
       Index->updatePreamble(testPath(Input.first()), "null", Ctx, PP,
-                            CanonIncludes);
+                            CanonIncludes, PI);
     });
     ParsedAST MainAST = TU.build();
     Index->updateMain(testPath(Input.first()), MainAST);
Index: clang-tools-extra/clangd/unittests/TestTU.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/TestTU.cpp
+++ clang-tools-extra/clangd/unittests/TestTU.cpp
@@ -164,9 +164,9 @@
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return std::get<0>(indexHeaderSymbols(/*Version=*/"null", AST.getASTContext(),
-                                        AST.getPreprocessor(),
-                                        AST.getCanonicalIncludes()));
+  return std::get<0>(indexHeaderSymbols(
+      /*Version=*/"null", AST.getASTContext(), AST.getPreprocessor(),
+      AST.getCanonicalIncludes(), AST.getPragmaIncludes()));
 }
 
 RefSlab TestTU::headerRefs() const {
@@ -179,7 +179,7 @@
   auto Idx = std::make_unique<FileIndex>();
   Idx->updatePreamble(testPath(Filename), /*Version=*/"null",
                       AST.getASTContext(), AST.getPreprocessor(),
-                      AST.getCanonicalIncludes());
+                      AST.getCanonicalIncludes(), AST.getPragmaIncludes());
   Idx->updateMain(testPath(Filename), AST);
   return std::move(Idx);
 }
Index: clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -18,6 +18,7 @@
 #include "TUScheduler.h"
 #include "TestFS.h"
 #include "TestIndex.h"
+#include "clang-include-cleaner/Record.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Path.h"
@@ -1131,7 +1132,8 @@
         : BlockVersion(BlockVersion), N(N) {}
     void onPreambleAST(PathRef Path, llvm::StringRef Version,
                        const CompilerInvocation &, ASTContext &Ctx,
-                       Preprocessor &, const CanonicalIncludes &) override {
+                       Preprocessor &, const CanonicalIncludes &,
+                       const include_cleaner::PragmaIncludes *) override {
       if (Version == BlockVersion)
         N.wait();
     }
@@ -1210,7 +1212,8 @@
 
     void onPreambleAST(PathRef Path, llvm::StringRef Version,
                        const CompilerInvocation &, ASTContext &Ctx,
-                       Preprocessor &, const CanonicalIncludes &) override {
+                       Preprocessor &, const CanonicalIncludes &,
+                       const include_cleaner::PragmaIncludes *) override {
       if (BuildBefore)
         ASSERT_TRUE(UnblockPreamble.wait(timeoutSeconds(5)))
             << "Expected notification";
@@ -1564,7 +1567,8 @@
         : Filenames(Filenames) {}
     void onPreambleAST(PathRef Path, llvm::StringRef Version,
                        const CompilerInvocation &CI, ASTContext &Ctx,
-                       Preprocessor &PP, const CanonicalIncludes &) override {
+                       Preprocessor &PP, const CanonicalIncludes &,
+                       const include_cleaner::PragmaIncludes *) override {
       // Deliberately no synchronization.
       // The PreambleThrottler should serialize these calls, if not then tsan
       // will find a bug here.
Index: clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
+++ clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
@@ -9,6 +9,8 @@
 #include "Annotations.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "URI.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/FileSystemOptions.h"
@@ -28,6 +30,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <utility>
 
 namespace clang {
 namespace clangd {
@@ -226,23 +229,21 @@
 
 class SymbolIndexActionFactory : public tooling::FrontendActionFactory {
 public:
-  SymbolIndexActionFactory(SymbolCollector::Options COpts,
-                           CommentHandler *PragmaHandler)
-      : COpts(std::move(COpts)), PragmaHandler(PragmaHandler) {}
+  SymbolIndexActionFactory(SymbolCollector::Options COpts)
+      : COpts(std::move(COpts)) {}
 
   std::unique_ptr<FrontendAction> create() override {
     class IndexAction : public ASTFrontendAction {
     public:
       IndexAction(std::shared_ptr<index::IndexDataConsumer> DataConsumer,
                   const index::IndexingOptions &Opts,
-                  CommentHandler *PragmaHandler)
+                  std::unique_ptr<include_cleaner::PragmaIncludes> PI)
           : DataConsumer(std::move(DataConsumer)), Opts(Opts),
-            PragmaHandler(PragmaHandler) {}
+            PI(std::move(PI)) {}
 
       std::unique_ptr<ASTConsumer>
       CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
-        if (PragmaHandler)
-          CI.getPreprocessor().addCommentHandler(PragmaHandler);
+        PI->record(CI);
         return createIndexingASTConsumer(DataConsumer, Opts,
                                          CI.getPreprocessorPtr());
       }
@@ -256,20 +257,22 @@
     private:
       std::shared_ptr<index::IndexDataConsumer> DataConsumer;
       index::IndexingOptions Opts;
-      CommentHandler *PragmaHandler;
+      std::unique_ptr<include_cleaner::PragmaIncludes> PI;
     };
     index::IndexingOptions IndexOpts;
     IndexOpts.SystemSymbolFilter =
         index::IndexingOptions::SystemSymbolFilterKind::All;
     IndexOpts.IndexFunctionLocals = true;
+    std::unique_ptr<include_cleaner::PragmaIncludes> PI =
+        std::make_unique<include_cleaner::PragmaIncludes>();
+    COpts.PragmaIncludes = PI.get();
     Collector = std::make_shared<SymbolCollector>(COpts);
     return std::make_unique<IndexAction>(Collector, std::move(IndexOpts),
-                                         PragmaHandler);
+                                         std::move(PI));
   }
 
   std::shared_ptr<SymbolCollector> Collector;
   SymbolCollector::Options COpts;
-  CommentHandler *PragmaHandler;
 };
 
 class SymbolCollectorTest : public ::testing::Test {
@@ -289,8 +292,7 @@
     llvm::IntrusiveRefCntPtr<FileManager> Files(
         new FileManager(FileSystemOptions(), InMemoryFileSystem));
 
-    auto Factory = std::make_unique<SymbolIndexActionFactory>(
-        CollectorOpts, PragmaHandler.get());
+    auto Factory = std::make_unique<SymbolIndexActionFactory>(CollectorOpts);
 
     std::vector<std::string> Args = {"symbol_collector", "-fsyntax-only",
                                      "-xc++", "-include", TestHeaderName};
@@ -324,7 +326,6 @@
   RefSlab Refs;
   RelationSlab Relations;
   SymbolCollector::Options CollectorOpts;
-  std::unique_ptr<CommentHandler> PragmaHandler;
 };
 
 TEST_F(SymbolCollectorTest, CollectSymbols) {
@@ -1549,7 +1550,7 @@
   auto Language = LangOptions();
   Language.CPlusPlus = true;
   Includes.addSystemHeadersMapping(Language);
-  CollectorOpts.Includes = &Includes;
+  CollectorOpts.SysHeaderMapping = &Includes;
   runSymbolCollector(
       R"cpp(
       namespace std {
@@ -1573,9 +1574,6 @@
 
 TEST_F(SymbolCollectorTest, IWYUPragma) {
   CollectorOpts.CollectIncludePath = true;
-  CanonicalIncludes Includes;
-  PragmaHandler = collectIWYUHeaderMaps(&Includes);
-  CollectorOpts.Includes = &Includes;
   const std::string Header = R"(
     // IWYU pragma: private, include the/good/header.h
     class Foo {};
@@ -1588,9 +1586,6 @@
 
 TEST_F(SymbolCollectorTest, IWYUPragmaWithDoubleQuotes) {
   CollectorOpts.CollectIncludePath = true;
-  CanonicalIncludes Includes;
-  PragmaHandler = collectIWYUHeaderMaps(&Includes);
-  CollectorOpts.Includes = &Includes;
   const std::string Header = R"(
     // IWYU pragma: private, include "the/good/header.h"
     class Foo {};
@@ -1601,6 +1596,27 @@
                                  includeHeader("\"the/good/header.h\""))));
 }
 
+TEST_F(SymbolCollectorTest, IWYUPragmaExport) {
+  CollectorOpts.CollectIncludePath = true;
+  const std::string Header = R"cpp(#pragma once
+    #include "exporter.h"
+  )cpp";
+  auto ExporterFile = testPath("exporter.h");
+  InMemoryFileSystem->addFile(
+      ExporterFile, 0, llvm::MemoryBuffer::getMemBuffer(R"cpp(#pragma once
+    #include "private.h" // IWYU pragma: export
+  )cpp"));
+  auto PrivateFile = testPath("private.h");
+  InMemoryFileSystem->addFile(
+      PrivateFile, 0, llvm::MemoryBuffer::getMemBuffer("class Foo {};"));
+  runSymbolCollector(Header, /*Main=*/"",
+                     /*ExtraArgs=*/{"-I", testRoot()});
+  EXPECT_THAT(Symbols, UnorderedElementsAre(AllOf(
+                           qName("Foo"),
+                           includeHeader(URI::create(ExporterFile).toString()),
+                           declURI(URI::create(PrivateFile).toString()))));
+}
+
 TEST_F(SymbolCollectorTest, SkipIncFileWhenCanonicalizeHeaders) {
   auto IncFile = testPath("test.inc");
   auto IncURI = URI::create(IncFile).toString();
@@ -1616,7 +1632,7 @@
   CanonicalIncludes Includes;
   Includes.addMapping(*File, "<canonical>");
   CollectorOpts.CollectIncludePath = true;
-  CollectorOpts.Includes = &Includes;
+  CollectorOpts.SysHeaderMapping = &Includes;
   runSymbolCollector(HeaderCode, /*Main=*/"",
                      /*ExtraArgs=*/{"-I", testRoot()});
   EXPECT_THAT(Symbols,
Index: clang-tools-extra/clangd/unittests/FileIndexTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/FileIndexTests.cpp
+++ clang-tools-extra/clangd/unittests/FileIndexTests.cpp
@@ -15,6 +15,7 @@
 #include "TestTU.h"
 #include "TestWorkspace.h"
 #include "URI.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
@@ -60,6 +61,11 @@
 MATCHER_P(numReferences, N, "") { return arg.References == N; }
 MATCHER_P(hasOrign, O, "") { return bool(arg.Origin & O); }
 
+MATCHER_P(includeHeader, P, "") {
+  return (arg.IncludeHeaders.size() == 1) &&
+         (arg.IncludeHeaders.begin()->IncludeHeader == P);
+}
+
 namespace clang {
 namespace clangd {
 namespace {
@@ -171,7 +177,7 @@
   auto AST = File.build();
   M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
                    AST.getASTContext(), AST.getPreprocessor(),
-                   AST.getCanonicalIncludes());
+                   AST.getCanonicalIncludes(), AST.getPragmaIncludes());
 }
 
 TEST(FileIndexTest, CustomizedURIScheme) {
@@ -234,6 +240,32 @@
               "<the/good/header.h>");
 }
 
+TEST(FileIndexTest, IWYUPragmaExport) {
+  FileIndex M;
+
+  TestTU File;
+  File.Code = R"cpp(#pragma once
+    #include "exporter.h"
+  )cpp";
+  File.HeaderFilename = "exporter.h";
+  File.HeaderCode = R"cpp(#pragma once
+    #include "private.h" // IWYU pragma: export
+  )cpp";
+  File.AdditionalFiles["private.h"] = "class Foo{};";
+  auto AST = File.build();
+  M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
+                   AST.getASTContext(), AST.getPreprocessor(),
+                   AST.getCanonicalIncludes(), AST.getPragmaIncludes());
+
+  auto Symbols = runFuzzyFind(M, "");
+  EXPECT_THAT(
+      Symbols,
+      UnorderedElementsAre(AllOf(
+          qName("Foo"),
+          includeHeader(URI::create(testPath(File.HeaderFilename)).toString()),
+          declURI(URI::create(testPath("private.h")).toString()))));
+}
+
 TEST(FileIndexTest, HasSystemHeaderMappingsInPreamble) {
   TestTU TU;
   TU.HeaderCode = "class Foo{};";
@@ -308,12 +340,13 @@
   buildPreamble(FooCpp, *CI, PI,
                 /*StoreInMemory=*/true,
                 [&](ASTContext &Ctx, Preprocessor &PP,
-                    const CanonicalIncludes &CanonIncludes) {
+                    const CanonicalIncludes &CanonIncludes,
+                    const include_cleaner::PragmaIncludes *PI) {
                   EXPECT_FALSE(IndexUpdated)
                       << "Expected only a single index update";
                   IndexUpdated = true;
                   Index.updatePreamble(FooCpp, /*Version=*/"null", Ctx, PP,
-                                       CanonIncludes);
+                                       CanonIncludes, PI);
                 });
   ASSERT_TRUE(IndexUpdated);
 
@@ -414,7 +447,7 @@
   FileIndex Index;
   Index.updatePreamble(testPath(TU.Filename), /*Version=*/"null",
                        AST.getASTContext(), AST.getPreprocessor(),
-                       AST.getCanonicalIncludes());
+                       AST.getCanonicalIncludes(), AST.getPragmaIncludes());
   SymbolID A = findSymbol(TU.headerSymbols(), "A").ID;
   uint32_t Results = 0;
   RelationsRequest Req;
@@ -535,7 +568,7 @@
   auto AST = File.build();
   M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
                    AST.getASTContext(), AST.getPreprocessor(),
-                   AST.getCanonicalIncludes());
+                   AST.getCanonicalIncludes(), AST.getPragmaIncludes());
   EXPECT_THAT(runFuzzyFind(M, ""), UnorderedElementsAre(qName("a")));
 
   File.Filename = "f2.cpp";
@@ -543,7 +576,7 @@
   AST = File.build();
   M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
                    AST.getASTContext(), AST.getPreprocessor(),
-                   AST.getCanonicalIncludes());
+                   AST.getCanonicalIncludes(), AST.getPragmaIncludes());
   EXPECT_THAT(runFuzzyFind(M, ""), UnorderedElementsAre(qName("b")));
 }
 
@@ -688,7 +721,7 @@
   auto AST = TestTU::withHeaderCode("int a;").build();
   FI.updateMain(FileName, AST);
   FI.updatePreamble(FileName, "v1", AST.getASTContext(), AST.getPreprocessor(),
-                    AST.getCanonicalIncludes());
+                    AST.getCanonicalIncludes(), AST.getPragmaIncludes());
 
   llvm::BumpPtrAllocator Alloc;
   MemoryTree MT(&Alloc);
Index: clang-tools-extra/clangd/unittests/CanonicalIncludesTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/CanonicalIncludesTests.cpp
+++ clang-tools-extra/clangd/unittests/CanonicalIncludesTests.cpp
@@ -30,39 +30,6 @@
   return *File;
 }
 
-TEST(CanonicalIncludesTest, CStandardLibrary) {
-  CanonicalIncludes CI;
-  auto Language = LangOptions();
-  Language.C11 = true;
-  CI.addSystemHeadersMapping(Language);
-  // Usual standard library symbols are mapped correctly.
-  EXPECT_EQ("<stdio.h>", CI.mapSymbol("", "printf", Language));
-  EXPECT_EQ("", CI.mapSymbol("", "unknown_symbol", Language));
-}
-
-TEST(CanonicalIncludesTest, CXXStandardLibrary) {
-  CanonicalIncludes CI;
-  auto Language = LangOptions();
-  Language.CPlusPlus = true;
-  CI.addSystemHeadersMapping(Language);
-
-  // Usual standard library symbols are mapped correctly.
-  EXPECT_EQ("<vector>", CI.mapSymbol("std::", "vector", Language));
-  EXPECT_EQ("<cstdio>", CI.mapSymbol("std::", "printf", Language));
-  // std::move is ambiguous, currently always mapped to <utility>
-  EXPECT_EQ("<utility>", CI.mapSymbol("std::", "move", Language));
-  EXPECT_EQ("<cstddef>", CI.mapSymbol("std::", "size_t", Language));
-  // Unknown std symbols aren't mapped.
-  EXPECT_EQ("", CI.mapSymbol("std::", "notathing", Language));
-  // iosfwd declares some symbols it doesn't own.
-  EXPECT_EQ("<ostream>", CI.mapSymbol("std::", "ostream", Language));
-  // And (for now) we assume it owns the others.
-  auto InMemFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
-  FileManager Files(FileSystemOptions(), InMemFS);
-  auto File = addFile(*InMemFS, Files, testPath("iosfwd"));
-  EXPECT_EQ("<iosfwd>", CI.mapHeader(File));
-}
-
 TEST(CanonicalIncludesTest, PathMapping) {
   auto InMemFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
   FileManager Files(FileSystemOptions(), InMemFS);
Index: clang-tools-extra/clangd/tool/Check.cpp
===================================================================
--- clang-tools-extra/clangd/tool/Check.cpp
+++ clang-tools-extra/clangd/tool/Check.cpp
@@ -47,6 +47,7 @@
 #include "SourceCode.h"
 #include "TidyProvider.h"
 #include "XRefs.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "refactor/Tweak.h"
@@ -230,12 +231,13 @@
     log("Building preamble...");
     Preamble = buildPreamble(File, *Invocation, Inputs, /*StoreInMemory=*/true,
                              [&](ASTContext &Ctx, Preprocessor &PP,
-                                 const CanonicalIncludes &Includes) {
+                                 const CanonicalIncludes &Includes,
+                                 const include_cleaner::PragmaIncludes *PI) {
                                if (!Opts.BuildDynamicSymbolIndex)
                                  return;
                                log("Indexing headers...");
                                Index.updatePreamble(File, /*Version=*/"null",
-                                                    Ctx, PP, Includes);
+                                                    Ctx, PP, Includes, PI);
                              });
     if (!Preamble) {
       elog("Failed to build preamble");
Index: clang-tools-extra/clangd/index/SymbolCollector.h
===================================================================
--- clang-tools-extra/clangd/index/SymbolCollector.h
+++ clang-tools-extra/clangd/index/SymbolCollector.h
@@ -9,6 +9,8 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOLCOLLECTOR_H
 
 #include "CollectMacros.h"
+#include "clang-include-cleaner/Record.h"
+#include "clang-include-cleaner/Types.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Ref.h"
 #include "index/Relation.h"
@@ -24,6 +26,7 @@
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 #include <functional>
 #include <optional>
 
@@ -56,9 +59,12 @@
     /// path.
     std::string FallbackDir;
     bool CollectIncludePath = false;
+    /// If set, this is used to map symbols to their #include paths
+    /// in system headers.
+    const CanonicalIncludes *SysHeaderMapping = nullptr;
     /// If set, this is used to map symbol #include path to a potentially
-    /// different #include path.
-    const CanonicalIncludes *Includes = nullptr;
+    /// different #include path specified by IWYU pragmas.
+    const include_cleaner::PragmaIncludes *PragmaIncludes = nullptr;
     // Populate the Symbol.References field.
     bool CountReferences = false;
     /// The symbol ref kinds that will be collected.
@@ -166,13 +172,25 @@
 
   // All Symbols collected from the AST.
   SymbolSlab::Builder Symbols;
-  // File IDs for Symbol.IncludeHeaders.
-  // The final spelling is calculated in finish().
+  // File IDs to distinguish imports from includes.
   llvm::DenseMap<SymbolID, FileID> IncludeFiles;
+  void setIncludeLocation(const Symbol &S, SourceLocation);
+
+  // Providers for Symbol.IncludeHeaders.
+  // The final spelling is calculated in finish().
+  llvm::DenseMap<SymbolID, llvm::SmallVector<include_cleaner::Header>>
+      SymbolProviders;
+  // Mapping from header object to final include spelling.
+  // Needs to be persisted here because Symbols don't own any of their data
+  // (i.e., header spellings in this case).
+  llvm::DenseMap<include_cleaner::Header, std::string> HeaderSpelling;
   // Files which contain ObjC symbols.
   // This is finalized and used in finish().
   llvm::DenseSet<FileID> FilesWithObjCConstructs;
-  void setIncludeLocation(const Symbol &S, SourceLocation);
+
+  void
+  setSymbolProviders(const Symbol &S,
+                     const llvm::SmallVector<include_cleaner::Header> Headers);
   // Indexed macros, to be erased if they turned out to be include guards.
   llvm::DenseSet<const IdentifierInfo *> IndexedMacros;
   // All refs collected from the AST. It includes:
@@ -184,6 +202,7 @@
   RelationSlab::Builder Relations;
   ASTContext *ASTCtx;
   Preprocessor *PP = nullptr;
+  const include_cleaner::PragmaIncludes *PI;
   std::shared_ptr<GlobalCodeCompletionAllocator> CompletionAllocator;
   std::unique_ptr<CodeCompletionTUInfo> CompletionTUInfo;
   Options Opts;
Index: clang-tools-extra/clangd/index/SymbolCollector.cpp
===================================================================
--- clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -13,8 +13,12 @@
 #include "ExpectedTypes.h"
 #include "SourceCode.h"
 #include "URI.h"
+#include "clang-include-cleaner/Analysis.h"
+#include "clang-include-cleaner/Record.h"
+#include "clang-include-cleaner/Types.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Relation.h"
+#include "index/Symbol.h"
 #include "index/SymbolID.h"
 #include "index/SymbolLocation.h"
 #include "clang/AST/Decl.h"
@@ -30,10 +34,16 @@
 #include "clang/Lex/Token.h"
 #include "clang/Tooling/Inclusions/HeaderAnalysis.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include <optional>
+#include <string>
+#include <utility>
 
 namespace clang {
 namespace clangd {
@@ -188,7 +198,7 @@
   // (IndexDataConsumer::setPreprocessor can happen before or after initialize)
   Preprocessor *&PP;
   const SourceManager &SM;
-  const CanonicalIncludes *Includes;
+  const CanonicalIncludes *SysHeaderMapping;
   llvm::StringRef FallbackDir;
   llvm::DenseMap<const FileEntry *, const std::string *> CacheFEToURI;
   llvm::StringMap<std::string> CachePathToURI;
@@ -196,12 +206,13 @@
   llvm::StringMap<std::string> CachePathToFrameworkSpelling;
   llvm::StringMap<FrameworkUmbrellaSpelling>
       CacheFrameworkToUmbrellaHeaderSpelling;
+  const include_cleaner::PragmaIncludes *PI;
 
 public:
   HeaderFileURICache(Preprocessor *&PP, const SourceManager &SM,
                      const SymbolCollector::Options &Opts)
-      : PP(PP), SM(SM), Includes(Opts.Includes), FallbackDir(Opts.FallbackDir) {
-  }
+      : PP(PP), SM(SM), SysHeaderMapping(Opts.SysHeaderMapping),
+        FallbackDir(Opts.FallbackDir), PI(Opts.PragmaIncludes) {}
 
   // Returns a canonical URI for the file \p FE.
   // We attempt to make the path absolute first.
@@ -234,6 +245,22 @@
     return R.first->second;
   }
 
+  // If a file is mapped by canonical headers, use that mapping, regardless
+  // of whether it's an otherwise-good header (header guards etc).
+  llvm::StringRef mapCanonical(FileID FID) {
+    if (SysHeaderMapping) {
+      llvm::StringRef Canonical =
+          SysHeaderMapping->mapHeader(*SM.getFileEntryRefForID(FID));
+      if (!Canonical.empty()) {
+        // If we had a mapping, always use it.
+        if (Canonical.startswith("<") || Canonical.startswith("\""))
+          return Canonical;
+        return toURI(Canonical);
+      }
+    }
+    return "";
+  }
+
 private:
   // This takes care of making paths absolute and path->URI caching, but no
   // FileManager-based canonicalization.
@@ -376,19 +403,19 @@
     const auto FE = SM.getFileEntryRefForID(FID);
     if (!FE || FE->getName().empty())
       return "";
-    llvm::StringRef Filename = FE->getName();
-    // If a file is mapped by canonical headers, use that mapping, regardless
-    // of whether it's an otherwise-good header (header guards etc).
-    if (Includes) {
-      llvm::StringRef Canonical =
-          Includes->mapHeader(*SM.getFileEntryRefForID(FID));
-      if (!Canonical.empty()) {
-        // If we had a mapping, always use it.
-        if (Canonical.startswith("<") || Canonical.startswith("\""))
-          return Canonical;
-        return toURI(Canonical);
-      }
-    }
+
+    // Fallback pragma handling for Objective-C.
+    const auto *FileEntry = SM.getFileEntryForID(FID);
+    for (const auto *Export : PI->getExporters(FileEntry, SM.getFileManager()))
+      return toURI(Export->tryGetRealPathName());
+
+    if (auto Verbatim = PI->getPublic(FileEntry); !Verbatim.empty())
+      return Verbatim;
+
+    auto Canonical = mapCanonical(FID);
+    if (!Canonical.empty())
+      return Canonical;
+
     // Framework headers are spelled as <FrameworkName/Foo.h>, not
     // "path/FrameworkName.framework/Headers/Foo.h".
     auto &HS = PP->getHeaderSearchInfo();
@@ -398,6 +425,7 @@
                 getFrameworkHeaderIncludeSpelling(*FE, HFI->Framework, HS))
           return *Spelling;
 
+    llvm::StringRef Filename = FE->getName();
     if (!tooling::isSelfContainedHeader(*FE, PP->getSourceManager(),
                                         PP->getHeaderSearchInfo())) {
       // A .inc or .def file is often included into a real header to define
@@ -430,7 +458,9 @@
   return Result;
 }
 
-SymbolCollector::SymbolCollector(Options Opts) : Opts(std::move(Opts)) {}
+SymbolCollector::SymbolCollector(Options Opts) : Opts(std::move(Opts)) {
+  this->PI = Opts.PragmaIncludes;
+}
 SymbolCollector::~SymbolCollector() = default;
 
 void SymbolCollector::initialize(ASTContext &Ctx) {
@@ -765,6 +795,11 @@
 
   IndexedMacros.insert(Name);
   setIncludeLocation(S, DefLoc);
+  llvm::SmallVector<include_cleaner::Header> Headers =
+      include_cleaner::headersForSymbol(
+          include_cleaner::Macro{const_cast<IdentifierInfo *>(Name), DefLoc},
+          SM, PI);
+  setSymbolProviders(S, Headers);
   Symbols.insert(S);
   return true;
 }
@@ -807,6 +842,13 @@
         PP->getSourceManager().getDecomposedExpansionLoc(Loc).first;
 }
 
+void SymbolCollector::setSymbolProviders(
+    const Symbol &S, const llvm::SmallVector<include_cleaner::Header> Headers) {
+  if (Opts.CollectIncludePath &&
+      shouldCollectIncludePath(S.SymInfo.Kind) != Symbol::Invalid)
+    SymbolProviders[S.ID] = std::move(Headers);
+}
+
 void SymbolCollector::finish() {
   // At the end of the TU, add 1 to the refcount of all referenced symbols.
   for (const auto &ID : ReferencedSymbols) {
@@ -834,50 +876,84 @@
   // We delay this until end of TU so header guards are all resolved.
   for (const auto &[SID, FID] : IncludeFiles) {
     if (const Symbol *S = Symbols.find(SID)) {
-      llvm::StringRef IncludeHeader;
-      // Look for an overridden include header for this symbol specifically.
-      if (Opts.Includes) {
-        IncludeHeader =
-            Opts.Includes->mapSymbol(S->Scope, S->Name, ASTCtx->getLangOpts());
-        if (!IncludeHeader.empty()) {
-          if (IncludeHeader.front() != '"' && IncludeHeader.front() != '<')
-            IncludeHeader = HeaderFileURIs->toURI(IncludeHeader);
-          else if (IncludeHeader == "<utility>" && S->Scope == "std::" &&
-                   S->Name == "move" && S->Signature.contains(','))
-            IncludeHeader = "<algorithm>";
-        }
+      // Determine if the FID is #include'd or #import'ed.
+      Symbol::IncludeDirective Directives = Symbol::Invalid;
+      auto CollectDirectives = shouldCollectIncludePath(S->SymInfo.Kind);
+      if ((CollectDirectives & Symbol::Include) != 0)
+        Directives |= Symbol::Include;
+      // Only allow #import for symbols from ObjC-like files.
+      if ((CollectDirectives & Symbol::Import) != 0) {
+        auto [It, Inserted] = FileToContainsImportsOrObjC.try_emplace(FID);
+        if (Inserted)
+          It->second = FilesWithObjCConstructs.contains(FID) ||
+                       tooling::codeContainsImports(
+                           ASTCtx->getSourceManager().getBufferData(FID));
+        if (It->second)
+          Directives |= Symbol::Import;
       }
-      // Otherwise find the approprate include header for the defining file.
-      if (IncludeHeader.empty())
-        IncludeHeader = HeaderFileURIs->getIncludeHeader(FID);
-
-      // Symbols in slabs aren't mutable, insert() has to walk all the strings
-      if (!IncludeHeader.empty()) {
-        Symbol::IncludeDirective Directives = Symbol::Invalid;
-        auto CollectDirectives = shouldCollectIncludePath(S->SymInfo.Kind);
-        if ((CollectDirectives & Symbol::Include) != 0)
-          Directives |= Symbol::Include;
-        // Only allow #import for symbols from ObjC-like files.
-        if ((CollectDirectives & Symbol::Import) != 0) {
-          auto [It, Inserted] = FileToContainsImportsOrObjC.try_emplace(FID);
-          if (Inserted)
-            It->second = FilesWithObjCConstructs.contains(FID) ||
-                         tooling::codeContainsImports(
-                             ASTCtx->getSourceManager().getBufferData(FID));
-          if (It->second)
-            Directives |= Symbol::Import;
-        }
-        if (Directives != Symbol::Invalid) {
-          Symbol NewSym = *S;
-          NewSym.IncludeHeaders.push_back({IncludeHeader, 1, Directives});
-          Symbols.insert(NewSym);
+
+      if (Directives == Symbol::Invalid)
+        continue;
+
+      Symbol NewSym;
+      NewSym = *S;
+      // Use the include location-based logic for Objective-C symbols.
+      if (Directives & Symbol::Import) {
+        // If not overridden manually, find the approprate include header for
+        // the defining file.
+        llvm::StringRef IncludeHeader{HeaderFileURIs->getIncludeHeader(FID)};
+        if (!IncludeHeader.empty())
+          NewSym.IncludeHeaders.push_back(
+              {HeaderFileURIs->getIncludeHeader(FID), 1, Directives});
+        Symbols.insert(NewSym);
+        continue;
+      }
+
+      // For #include's, use the providers computed by the include-cleaner
+      // library.
+      if (Directives == Symbol::Include) {
+        auto It = SymbolProviders.find(SID);
+        if (It != SymbolProviders.end()) {
+          for (const auto &H : It->second) {
+            const auto [SpellingIt, Inserted] = HeaderSpelling.try_emplace(H);
+            if (Inserted) {
+              switch (H.kind()) {
+              case include_cleaner::Header::Kind::Verbatim:
+                SpellingIt->second = H.verbatim().str();
+                break;
+              case include_cleaner::Header::Kind::Standard:
+                SpellingIt->second = H.standard().name().str();
+                break;
+              case include_cleaner::Header::Kind::Physical:
+                auto &SM = ASTCtx->getSourceManager();
+                auto Canonical =
+                    HeaderFileURIs->mapCanonical(SM.getOrCreateFileID(
+                        H.physical(), SrcMgr::CharacteristicKind::C_User));
+                if (!Canonical.empty())
+                  SpellingIt->second = Canonical;
+                if (!tooling::isSelfContainedHeader(H.physical(),
+                                                    PP->getSourceManager(),
+                                                    PP->getHeaderSearchInfo()))
+                  break;
+                if (SpellingIt->second.empty())
+                  SpellingIt->second =
+                      HeaderFileURIs->toURI(H.physical()->tryGetRealPathName());
+                break;
+              }
+            }
+            if (!SpellingIt->second.empty())
+              NewSym.IncludeHeaders.push_back(
+                  {SpellingIt->second, 1, Directives});
+          }
         }
+        Symbols.insert(NewSym);
       }
     }
   }
 
   ReferencedSymbols.clear();
   IncludeFiles.clear();
+  SymbolProviders.clear();
   FilesWithObjCConstructs.clear();
 }
 
@@ -952,6 +1028,9 @@
 
   Symbols.insert(S);
   setIncludeLocation(S, ND.getLocation());
+  llvm::SmallVector<include_cleaner::Header> Headers =
+      include_cleaner::headersForSymbol(include_cleaner::Symbol{ND}, SM, PI);
+  setSymbolProviders(S, Headers);
   if (S.SymInfo.Lang == index::SymbolLanguage::ObjC)
     FilesWithObjCConstructs.insert(FID);
   return Symbols.find(S.ID);
Index: clang-tools-extra/clangd/index/IndexAction.cpp
===================================================================
--- clang-tools-extra/clangd/index/IndexAction.cpp
+++ clang-tools-extra/clangd/index/IndexAction.cpp
@@ -9,7 +9,9 @@
 #include "IndexAction.h"
 #include "AST.h"
 #include "Headers.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/Relation.h"
+#include "index/SymbolCollector.h"
 #include "index/SymbolOrigin.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
@@ -127,6 +129,7 @@
 public:
   IndexAction(std::shared_ptr<SymbolCollector> C,
               std::unique_ptr<CanonicalIncludes> Includes,
+              std::unique_ptr<include_cleaner::PragmaIncludes> PI,
               const index::IndexingOptions &Opts,
               std::function<void(SymbolSlab)> SymbolsCallback,
               std::function<void(RefSlab)> RefsCallback,
@@ -135,8 +138,7 @@
       : SymbolsCallback(SymbolsCallback), RefsCallback(RefsCallback),
         RelationsCallback(RelationsCallback),
         IncludeGraphCallback(IncludeGraphCallback), Collector(C),
-        Includes(std::move(Includes)), Opts(Opts),
-        PragmaHandler(collectIWYUHeaderMaps(this->Includes.get())) {
+        Includes(std::move(Includes)), PI(std::move(PI)), Opts(Opts) {
     this->Opts.ShouldTraverseDecl = [this](const Decl *D) {
       // Many operations performed during indexing is linear in terms of depth
       // of the decl (USR generation, name lookups, figuring out role of a
@@ -154,8 +156,8 @@
 
   std::unique_ptr<ASTConsumer>
   CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
-    CI.getPreprocessor().addCommentHandler(PragmaHandler.get());
     Includes->addSystemHeadersMapping(CI.getLangOpts());
+    PI->record(CI.getPreprocessor());
     if (IncludeGraphCallback != nullptr)
       CI.getPreprocessor().addPPCallbacks(
           std::make_unique<IncludeGraphCollector>(CI.getSourceManager(), IG));
@@ -202,8 +204,8 @@
   std::function<void(IncludeGraph)> IncludeGraphCallback;
   std::shared_ptr<SymbolCollector> Collector;
   std::unique_ptr<CanonicalIncludes> Includes;
+  std::unique_ptr<include_cleaner::PragmaIncludes> PI;
   index::IndexingOptions Opts;
-  std::unique_ptr<CommentHandler> PragmaHandler;
   IncludeGraph IG;
 };
 
@@ -229,11 +231,13 @@
     Opts.RefsInHeaders = true;
   }
   auto Includes = std::make_unique<CanonicalIncludes>();
-  Opts.Includes = Includes.get();
+  Opts.SysHeaderMapping = Includes.get();
+  auto PragmaIncludes = std::make_unique<include_cleaner::PragmaIncludes>();
+  Opts.PragmaIncludes = PragmaIncludes.get();
   return std::make_unique<IndexAction>(
-      std::make_shared<SymbolCollector>(std::move(Opts)), std::move(Includes),
-      IndexOpts, SymbolsCallback, RefsCallback, RelationsCallback,
-      IncludeGraphCallback);
+      std::make_shared<SymbolCollector>(Opts), std::move(Includes),
+      std::move(PragmaIncludes), IndexOpts, SymbolsCallback, RefsCallback,
+      RelationsCallback, IncludeGraphCallback);
 }
 
 } // namespace clangd
Index: clang-tools-extra/clangd/index/FileIndex.h
===================================================================
--- clang-tools-extra/clangd/index/FileIndex.h
+++ clang-tools-extra/clangd/index/FileIndex.h
@@ -16,6 +16,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "Headers.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "index/Merge.h"
@@ -113,7 +114,8 @@
   /// Update preamble symbols of file \p Path with all declarations in \p AST
   /// and macros in \p PP.
   void updatePreamble(PathRef Path, llvm::StringRef Version, ASTContext &AST,
-                      Preprocessor &PP, const CanonicalIncludes &Includes);
+                      Preprocessor &PP, const CanonicalIncludes &Includes,
+                      const include_cleaner::PragmaIncludes *PI);
   void updatePreamble(IndexFileIn);
 
   /// Update symbols and references from main file \p Path with
@@ -162,7 +164,8 @@
 /// included headers.
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
                              Preprocessor &PP,
-                             const CanonicalIncludes &Includes);
+                             const CanonicalIncludes &Includes,
+                             const include_cleaner::PragmaIncludes *PI);
 
 /// Takes slabs coming from a TU (multiple files) and shards them per
 /// declaration location.
Index: clang-tools-extra/clangd/index/FileIndex.cpp
===================================================================
--- clang-tools-extra/clangd/index/FileIndex.cpp
+++ clang-tools-extra/clangd/index/FileIndex.cpp
@@ -9,6 +9,7 @@
 #include "FileIndex.h"
 #include "CollectMacros.h"
 #include "ParsedAST.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
@@ -46,11 +47,14 @@
 SlabTuple indexSymbols(ASTContext &AST, Preprocessor &PP,
                        llvm::ArrayRef<Decl *> DeclsToIndex,
                        const MainFileMacros *MacroRefsToIndex,
-                       const CanonicalIncludes &Includes, bool IsIndexMainAST,
-                       llvm::StringRef Version, bool CollectMainFileRefs) {
+                       const CanonicalIncludes &Includes,
+                       const include_cleaner::PragmaIncludes *PI,
+                       bool IsIndexMainAST, llvm::StringRef Version,
+                       bool CollectMainFileRefs) {
   SymbolCollector::Options CollectorOpts;
   CollectorOpts.CollectIncludePath = true;
-  CollectorOpts.Includes = &Includes;
+  CollectorOpts.SysHeaderMapping = &Includes;
+  CollectorOpts.PragmaIncludes = PI;
   CollectorOpts.CountReferences = false;
   CollectorOpts.Origin =
       IsIndexMainAST ? SymbolOrigin::Open : SymbolOrigin::Preamble;
@@ -222,18 +226,19 @@
 SlabTuple indexMainDecls(ParsedAST &AST) {
   return indexSymbols(
       AST.getASTContext(), AST.getPreprocessor(), AST.getLocalTopLevelDecls(),
-      &AST.getMacros(), AST.getCanonicalIncludes(),
+      &AST.getMacros(), AST.getCanonicalIncludes(), AST.getPragmaIncludes(),
       /*IsIndexMainAST=*/true, AST.version(), /*CollectMainFileRefs=*/true);
 }
 
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
                              Preprocessor &PP,
-                             const CanonicalIncludes &Includes) {
+                             const CanonicalIncludes &Includes,
+                             const include_cleaner::PragmaIncludes *PI) {
   std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
   return indexSymbols(AST, PP, DeclsToIndex,
-                      /*MainFileMacros=*/nullptr, Includes,
+                      /*MainFileMacros=*/nullptr, Includes, PI,
                       /*IsIndexMainAST=*/false, Version,
                       /*CollectMainFileRefs=*/false);
 }
@@ -458,10 +463,11 @@
 
 void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
                                ASTContext &AST, Preprocessor &PP,
-                               const CanonicalIncludes &Includes) {
+                               const CanonicalIncludes &Includes,
+                               const include_cleaner::PragmaIncludes *PI) {
   IndexFileIn IF;
   std::tie(IF.Symbols, std::ignore, IF.Relations) =
-      indexHeaderSymbols(Version, AST, PP, Includes);
+      indexHeaderSymbols(Version, AST, PP, Includes, PI);
   updatePreamble(std::move(IF));
 }
 
Index: clang-tools-extra/clangd/index/CanonicalIncludes.h
===================================================================
--- clang-tools-extra/clangd/index/CanonicalIncludes.h
+++ clang-tools-extra/clangd/index/CanonicalIncludes.h
@@ -39,13 +39,6 @@
   /// Adds a file-to-string mapping from \p ID to \p CanonicalPath.
   void addMapping(FileEntryRef Header, llvm::StringRef CanonicalPath);
 
-  /// Returns the overridden include for a qualified symbol with, or "".
-  /// \p Scope and \p Name concatenation forms the fully qualified name.
-  /// \p Scope is the qualifier with the trailing "::" (e.g. "std::") or empty
-  /// (for global namespace).
-  llvm::StringRef mapSymbol(llvm::StringRef Scope, llvm::StringRef Name,
-                            const LangOptions &L) const;
-
   /// Returns the overridden include for files in \p Header, or "".
   llvm::StringRef mapHeader(FileEntryRef Header) const;
 
@@ -66,31 +59,6 @@
   /// Used only for mapping standard headers.
   const llvm::StringMap<llvm::StringRef> *StdSuffixHeaderMapping = nullptr;
 };
-
-/// Returns a CommentHandler that parses pragma comment on include files to
-/// determine when we should include a different header from the header that
-/// directly defines a symbol. Mappinps are registered with \p Includes.
-///
-/// Currently it only supports IWYU private pragma:
-/// https://github.com/include-what-you-use/include-what-you-use/blob/master/docs/IWYUPragmas.md#iwyu-pragma-private
-///
-/// We ignore other pragmas:
-/// - keep: this is common but irrelevant: we do not currently remove includes
-/// - export: this is common and potentially interesting, there are three cases:
-///    * Points to a public header (common): we can suppress include2 if you
-///      already have include1. Only marginally useful.
-///    * Points to a private header annotated with `private` (somewhat common):
-///      Not incrementally useful as we support private.
-///    * Points to a private header without pragmas (rare). This is a reversed
-///      private pragma, and is valuable but too rare to be worthwhile.
-/// - no_include: this is about as common as private, but only affects the
-///   current file, so the value is smaller. We could add support.
-/// - friend: this is less common than private, has implementation difficulties,
-///   and affects behavior in a limited scope.
-/// - associated: extremely rare
-std::unique_ptr<CommentHandler>
-collectIWYUHeaderMaps(CanonicalIncludes *Includes);
-
 } // namespace clangd
 } // namespace clang
 
Index: clang-tools-extra/clangd/index/CanonicalIncludes.cpp
===================================================================
--- clang-tools-extra/clangd/index/CanonicalIncludes.cpp
+++ clang-tools-extra/clangd/index/CanonicalIncludes.cpp
@@ -7,14 +7,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "CanonicalIncludes.h"
-#include "Headers.h"
 #include "clang/Basic/FileEntry.h"
-#include "clang/Tooling/Inclusions/HeaderAnalysis.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/Support/FileSystem/UniqueID.h"
 #include "llvm/Support/Path.h"
-#include <algorithm>
 
 namespace clang {
 namespace clangd {
@@ -701,54 +697,6 @@
   return "";
 }
 
-llvm::StringRef CanonicalIncludes::mapSymbol(llvm::StringRef Scope,
-                                             llvm::StringRef Name,
-                                             const LangOptions &L) const {
-  tooling::stdlib::Lang Lang;
-  if (L.CPlusPlus)
-    Lang = tooling::stdlib::Lang::CXX;
-  else if (L.C11)
-    Lang = tooling::stdlib::Lang::C;
-  else
-    return "";
-  // FIXME: remove the following special cases when the tooling stdlib supports
-  // them.
-  // There are two std::move()s, this is by far the most common.
-  if (Scope == "std::" && Name == "move")
-    return "<utility>";
-  if (auto StdSym = tooling::stdlib::Symbol::named(Scope, Name, Lang))
-    if (auto Header = StdSym->header())
-      return Header->name();
-  return "";
-}
-
-std::unique_ptr<CommentHandler>
-collectIWYUHeaderMaps(CanonicalIncludes *Includes) {
-  class PragmaCommentHandler : public clang::CommentHandler {
-  public:
-    PragmaCommentHandler(CanonicalIncludes *Includes) : Includes(Includes) {}
-
-    bool HandleComment(Preprocessor &PP, SourceRange Range) override {
-      auto Pragma = tooling::parseIWYUPragma(
-          PP.getSourceManager().getCharacterData(Range.getBegin()));
-      if (!Pragma || !Pragma->consume_front("private, include "))
-        return false;
-      auto &SM = PP.getSourceManager();
-      // We always insert using the spelling from the pragma.
-      if (auto *FE = SM.getFileEntryForID(SM.getFileID(Range.getBegin())))
-        Includes->addMapping(FE->getLastRef(),
-                             isLiteralInclude(*Pragma)
-                                 ? Pragma->str()
-                                 : ("\"" + *Pragma + "\"").str());
-      return false;
-    }
-
-  private:
-    CanonicalIncludes *const Includes;
-  };
-  return std::make_unique<PragmaCommentHandler>(Includes);
-}
-
 void CanonicalIncludes::addSystemHeadersMapping(const LangOptions &Language) {
   // FIXME: remove the std header mapping once we support ambiguous symbols, now
   // it serves as a fallback to disambiguate:
Index: clang-tools-extra/clangd/TUScheduler.h
===================================================================
--- clang-tools-extra/clangd/TUScheduler.h
+++ clang-tools-extra/clangd/TUScheduler.h
@@ -13,6 +13,7 @@
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Function.h"
 #include "support/MemoryTree.h"
@@ -163,7 +164,8 @@
   /// file. AST node in the current file should be observed on onMainAST call.
   virtual void onPreambleAST(PathRef Path, llvm::StringRef Version,
                              const CompilerInvocation &CI, ASTContext &Ctx,
-                             Preprocessor &PP, const CanonicalIncludes &) {}
+                             Preprocessor &PP, const CanonicalIncludes &,
+                             const include_cleaner::PragmaIncludes *) {}
 
   /// The argument function is run under the critical section guarding against
   /// races when closing the files.
Index: clang-tools-extra/clangd/TUScheduler.cpp
===================================================================
--- clang-tools-extra/clangd/TUScheduler.cpp
+++ clang-tools-extra/clangd/TUScheduler.cpp
@@ -54,6 +54,7 @@
 #include "GlobalCompilationDatabase.h"
 #include "ParsedAST.h"
 #include "Preamble.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
@@ -1081,9 +1082,10 @@
   LatestBuild = clang::clangd::buildPreamble(
       FileName, *Req.CI, Inputs, StoreInMemory,
       [&](ASTContext &Ctx, Preprocessor &PP,
-          const CanonicalIncludes &CanonIncludes) {
+          const CanonicalIncludes &CanonIncludes,
+          const include_cleaner::PragmaIncludes *PI) {
         Callbacks.onPreambleAST(FileName, Inputs.Version, *Req.CI, Ctx, PP,
-                                CanonIncludes);
+                                CanonIncludes, PI);
       },
       &Stats);
   if (!LatestBuild)
Index: clang-tools-extra/clangd/Preamble.h
===================================================================
--- clang-tools-extra/clangd/Preamble.h
+++ clang-tools-extra/clangd/Preamble.h
@@ -80,8 +80,9 @@
   bool MainIsIncludeGuarded = false;
 };
 
-using PreambleParsedCallback = std::function<void(ASTContext &, Preprocessor &,
-                                                  const CanonicalIncludes &)>;
+using PreambleParsedCallback =
+    std::function<void(ASTContext &, Preprocessor &, const CanonicalIncludes &,
+                       const include_cleaner::PragmaIncludes *)>;
 
 /// Timings and statistics from the premble build. Unlike PreambleData, these
 /// do not need to be stored for later, but can be useful for logging, metrics,
Index: clang-tools-extra/clangd/Preamble.cpp
===================================================================
--- clang-tools-extra/clangd/Preamble.cpp
+++ clang-tools-extra/clangd/Preamble.cpp
@@ -100,7 +100,8 @@
   void AfterExecute(CompilerInstance &CI) override {
     if (ParsedCallback) {
       trace::Span Tracer("Running PreambleCallback");
-      ParsedCallback(CI.getASTContext(), CI.getPreprocessor(), CanonIncludes);
+      ParsedCallback(CI.getASTContext(), CI.getPreprocessor(), CanonIncludes,
+                     &Pragmas);
     }
 
     const SourceManager &SM = CI.getSourceManager();
@@ -149,11 +150,6 @@
         collectPragmaMarksCallback(*SourceMgr, Marks));
   }
 
-  CommentHandler *getCommentHandler() override {
-    IWYUHandler = collectIWYUHeaderMaps(&CanonIncludes);
-    return IWYUHandler.get();
-  }
-
   static bool isLikelyForwardingFunction(FunctionTemplateDecl *FT) {
     const auto *FD = FT->getTemplatedDecl();
     const auto NumParams = FD->getNumParams();
@@ -209,7 +205,6 @@
   MainFileMacros Macros;
   std::vector<PragmaMark> Marks;
   bool IsMainFileIncludeGuarded = false;
-  std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   const clang::LangOptions *LangOpts = nullptr;
   const SourceManager *SourceMgr = nullptr;
   const Preprocessor *PP = nullptr;
Index: clang-tools-extra/clangd/ParsedAST.cpp
===================================================================
--- clang-tools-extra/clangd/ParsedAST.cpp
+++ clang-tools-extra/clangd/ParsedAST.cpp
@@ -625,9 +625,6 @@
     CanonIncludes = Preamble->CanonIncludes;
   else
     CanonIncludes.addSystemHeadersMapping(Clang->getLangOpts());
-  std::unique_ptr<CommentHandler> IWYUHandler =
-      collectIWYUHeaderMaps(&CanonIncludes);
-  PP.addCommentHandler(IWYUHandler.get());
 
   // Collect tokens of the main file.
   syntax::TokenCollector CollectTokens(PP);
Index: clang-tools-extra/clangd/ClangdServer.cpp
===================================================================
--- clang-tools-extra/clangd/ClangdServer.cpp
+++ clang-tools-extra/clangd/ClangdServer.cpp
@@ -23,6 +23,7 @@
 #include "SourceCode.h"
 #include "TUScheduler.h"
 #include "XRefs.h"
+#include "clang-include-cleaner/Record.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Merge.h"
@@ -69,15 +70,15 @@
 
   void onPreambleAST(PathRef Path, llvm::StringRef Version,
                      const CompilerInvocation &CI, ASTContext &Ctx,
-                     Preprocessor &PP,
-                     const CanonicalIncludes &CanonIncludes) override {
+                     Preprocessor &PP, const CanonicalIncludes &CanonIncludes,
+                     const include_cleaner::PragmaIncludes *PI) override {
     // If this preamble uses a standard library we haven't seen yet, index it.
     if (FIndex)
       if (auto Loc = Stdlib->add(*CI.getLangOpts(), PP.getHeaderSearchInfo()))
         indexStdlib(CI, std::move(*Loc));
 
     if (FIndex)
-      FIndex->updatePreamble(Path, Version, Ctx, PP, CanonIncludes);
+      FIndex->updatePreamble(Path, Version, Ctx, PP, CanonIncludes, PI);
   }
 
   void indexStdlib(const CompilerInvocation &CI, StdLibLocation Loc) {
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to