[clangd] Update symbol collector to use include-cleaner.

[VitaNuo]

Remove redundant variables.

[VitaNuo]

Remove redundancy.

[kadircet]

clang-tools-extra/clangd/index/SymbolCollector.cpp

393–395

sorry i don't understand what this logic is trying to achieve.

we seem to be prefering "first" exporting header, over the headers that directly provide the symbol. Also comment says it's done for objc, but there's nothing limiting this logic to ObjC. any chance this is unintended?

791

s/ASTCtx->getSourceManager()/SM

903

this is using legacy mappings for non-objc symbols too

915

getIncludeHeader might return an empty string

928

you're iterating over multiple headers, and only keeping references for the last one alive.

instead of storing these in the class state, you can have a llvm::DenseMap<Header, std::string> HeaderToSpelling; inside this function and later on just do:

for(const auto& H: It->second) {
  auto [SpellingIt, Inserted] = HeaderToSpelling.try_emplace(H);
  if(Inserted) {
      switch (...) { /* Update *SpellingIt */ }
  }
  NewSym.IncludeHeaders.push_back(*SpellingIt, 1, Directives);

934–936

we actually only want to use HeaderFileURIs->toURI here and not the getIncludeHeader, because:

• We want to make use of mapping logic in include-cleaner solely, and we already have all that information in Providers mapping.
• The extra logic we want to apply for ObjC symbols are already applied before reaching here.

the way we create a FileID here is a little bit iffy, by using toURI, hopefully we can avoid that conversion.

the only remaining issue is, we actually still want to use system header mappings inside CanonicalIncludes until we have better support for them in include-cleaner library itself. So I think we should still call Includes->mapHeader(H.physical()) before using toURI.

clang-tools-extra/clangd/index/SymbolCollector.h

131

unfortunately this alternative is called after we've parsed the file, from clangd/index/FileIndex.cpp, indexSymbols. hence attaching PI to PPCallbacks won't have any effect.

This is working unintentionally because we're still populating CanonicalIncludes with legacy IWYUPragmaHandler in clangd, and using that to perform private -> public mappings.

We should instead propagate the PragmaIncludes we have in ParsedAST and in preamble builds here. You should be able to test the new behaviour clang-tools-extra/clangd/unittests/FileIndexTests.cpp to make sure pragma handling through dynamic indexing code paths keep working as expected using an export pragma.

[kadircet]

clang-tools-extra/clangd/index/SymbolCollector.cpp

837

you can take this by value and std::move into the map

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments!

clang-tools-extra/clangd/index/SymbolCollector.cpp

393–395

No, this was intended, but possibly I didn't get it right.

comment says it's done for objc, but there's nothing limiting this logic to objc

AFAIU this whole code path will now only be reachable for objc symbols. I have removed pragma includes from the canonical include mapping now, leaving the pragma handling to include cleaner. However, that only triggers for C/C++, so in case we need the export and private pragmas for objc, we need to re-insert the handling for them somewhere.
There is no pre-existing test case for pragmas in objc and there is no usage of these pragmas for objc code in the code base either, so I have no way of knowing if this is actually needed. But I believe you've mentioned in some discussion we should still handle the pragmas for objc.

we seem to be preferring "first" exporting header, over the headers that directly provide the symbol.

Isn't that correct if there's an IWYU export pragma involved? The snippet comes from include_cleaner::findHeaders, with the only difference that it stops at the first exporter (since the canonical include mapping also just stored one mapping for the header).

Let me know how to do it better, or maybe if this is necessary at all. Honestly, I am not sure about this, since, as mentioned, there are no export or private/public pragmas in objc files in the codebase atm.

903

Removed this whole passage.

934–936

Ok SGTM. It seems, though, that we still need the "iffy" FID for the canonical mapping.. Let me know if you know a better way.

clang-tools-extra/clangd/index/SymbolCollector.h

131

Ok thank you. I'm using PragmaIncludes from AST and preamble builds for the dynamic index now.

For the background index, I've finally managed to move pragma recording to the IndexAction now.

I've also removed the redundant (to include cleaner) comment handlers from the AST build and preamble building logic.

Also removed the mapSymbol method from CanonicalIncludes, since it had one usage which should now be covered by the include cleaner, I believe.

[kadircet]

clang-tools-extra/clangd/ParsedAST.cpp

628 ↗

(On Diff #532964)

can you also add a FIXME here saying that we should attach PragmaInclude callbacks to main file ast build too?

this is not a recent regression as we were not doing it properly in previous version either. IWYU handler we attach only cares about private pragmas, which don't mean much inside the main file. but we actually need it for other pragmas like keep/export/no_include etc. (no need to do it in this patch as it's already getting quite big, let's do that as a follow up)

clang-tools-extra/clangd/Preamble.h

83 ↗

(On Diff #532964)

you need to rebase as we had some big changes around this logic recently. we should pass the PragmaIncludes also as a shared_ptr (and store as one inside the preambledata)

clang-tools-extra/clangd/index/CanonicalIncludes.h

40 ↗

(On Diff #532964)

AFAICT, only users of this endpoint were in IWYUCommentHandler (and some tests). we can also get rid of this one now.

More importantly now this turns into a mapper used only by symbolcollector, that can be cheaply created at use time (we just copy around a pointer every time we want to create it). so you can drop CanonicalIncludes from all of the interfaces, including SymbolCollector::Options and create one on demand in HeaderFileURICache. all we need is langopts, and it's available through preprocessor (not necessarily on construction, but at the time when we want to do the mappings).

As you're already touching all of the interfaces that propagate CanonicalIncludes around, hopefully this change should only make things simpler (and you already need to touch them when you rebase), but if that feels like too much churn in this patch feel free to do that in a follow up as well.

43 ↗

(On Diff #532964)

let's mention that this will always return a verbatim spelling, e.g. with quotes or angles.

57 ↗

(On Diff #532964)

you can drop this one too, only addMapping would introduce mappings here, but that's going away too.

clang-tools-extra/clangd/index/SymbolCollector.cpp

250

let's change the interface here to take in a FileEntry instead.

251

nit: early exits

if (!SysHeaderMapping)
  return "";
auto Canonical = ...;
if(Canonical.empty())
  return "";
...

258

we should never hit this code path anymore. so feel free to convert the if above to an assert

393–395

existing IWYUCommentHandler only handled private pragmas, so export pragmas were already dropped on the floor, we shouldn't introduce handling for them in here now. the mappings you perform for private/public pragmas and system headers below is enough.

(P.S. using exporters as "alternatives" vs "sole provider" are quite different, policy in include_cleaner does the former whereas logic in here does the latter)

417–418

let's move this logic into mapCanonical too.

420

nit: if (auto Canonical = ...; !Canonical.empty()) return Canonical;

836

what about merging this one and setIncludeLocation ?

e.g:

void SymbolCollector::setIncludeLocation(const Symbol &IdxS, SourceLocation DefLoc, const include_cleaner::Symbol &Sym) {
if (!Opts.CollectIncludePath ||
      shouldCollectIncludePath(S.SymInfo.Kind) == Symbol::Invalid)
  return;
IncludeFiles[SID] = ...;
SymbolProviders[SID] = headersForSymbol(...);
}

839

because you're taking in Headers as const, this move is not actually moving.

868

let's iterate over SymbolProviders instead, as IncludeFiles is a legacy struct now.

869–880

inside of this branch is getting crowded, maybe early exit instead?

const Symbol *S = Symbols.find(SID);
if(!S) continue;
...

919

no need to call insert if we didn't modify IncludeHeaders, so maybe reflow to something like:

if (Directives & Import) {
  if(auto IncHeader = ..; !IncHeader.empty()) {
     auto NewSym = *S;
     NewSym.IncludeHeaders....;
     Symbols.insert(NewSym);
  }
  // FIXME: Use providers from include-cleaner once it's polished for ObjC too.
  continue;
}

925

s/if/assert

941

no need for the getOrCreateFileID after you change parameter for mapCanonical to be a file entry.

943

nit:

if(auto Canonical = mapCanonical(H.physical()); !Canonical.empty())
  SpellingIt->second = Canonical;

945

this should be else if. if we performed a mapping, we no longer care about self-containedness of original header.

946

s/PP->getSourceManager()/SM

949

s/if (SpellingIt->second.empty())/else

clang-tools-extra/clangd/index/SymbolCollector.h

176

we seem to be populating this map for all kinds of symbols, as long as shouldCollectIncludePath returns non-invalid. what's up with distinguishing imports from includes?

187

no need to have this in the class state, you can just inline this into ::finish method instead.

192

let's move this closer to the SymbolProviders map

205

we already store this in Opts, can we just use it from there instead of creating an extra name for it?

[kadircet]

clang-tools-extra/clangd/index/SymbolCollector.cpp

836

also sorry for missing this so far, but i think we should limit this to first provider initially. as otherwise we'll have providers that just accidentally declare certain symbols and we don't want them to be inserted (unless they're the only provider and end up at the top) just because our ranking heuristics fail or something.

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for comments!

clang-tools-extra/clangd/index/CanonicalIncludes.h

40 ↗

(On Diff #532964)

... and create one on demand in HeaderFileURICache. all we need is langopts, and it's available through preprocessor (not necessarily on construction, but at the time when we want to do the mappings

This time sounds a little late, since we don't want to rebuild the system header mapping every time we request a mapping for a particular header.
Cache construction time doesn't work, since it happens before the preprocessor is set in the symbol collector (and we need the preprocessor to retrieve language options).
So far the safe place I've found is right after the preprocessor is set in the symbol collector. Another option is to collect the system header options in the finish() method, but I can't see why we would need to delay until then.

clang-tools-extra/clangd/index/SymbolCollector.cpp

250

I believe it has to be FileEntryRef instead. This is what the CanonicalIncludes::mapHeader expects, and this is also seems right (somewhere in the Clang docs I've read "everything that needs a name should use FileEntryRef", and name is exactly what CanonicalIncludes::mapHeader needs).

For the case of a physical header (i.e., FileEntry instance) I've switched to using getLastRef now. There's also another API FileManager::getFileRef where I could try passing the result of tryGetRealPathName. I am not sure I have enough information to judge if any of these options is distinctly better.

417–418

This would deliver wrong results for C++. We would basically get every IWYU-public header twice: once through the verbatim branch of the include_cleaner results, and then once again when mapping the physical private file to a canonical version (e.g., for system headers).

839

Ok, this is not relevant anymore, but I am not sure I understand why. Am I not allowed to say that I don't need this variable/memory anymore, even though it's const?

868

Not sure I'm following. Iterating over SymbolProviders means retrieving include_cleaner::Headers. How would I handle the entire Obj-C part then, without the FID of the include header?

941

Still need the getLastRef or FileManager::getFileRef, if I understand correctly.

clang-tools-extra/clangd/index/SymbolCollector.h

176

Yeah that's right. I've been trying to express that (after populating this map) we are not using this FID to compute the IncludeHeaders anymore (only for the ObjC case), but rather use it to determine if the Directives should contain Symbol::Import (ie., if this FID contains Obj-C constructs/imports).

I'll try to be a bit more verbose and precise.

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

85 ↗

(On Diff #535740)

nit: PI(std::move(PI))

clang-tools-extra/clangd/TUScheduler.cpp

1086 ↗

(On Diff #535740)

nit: s/PI/std::move(PI)

clang-tools-extra/clangd/index/CanonicalIncludes.cpp

672 ↗

(On Diff #535740)

we're no longer using any FileEntry pieces of this parameter you can take in just a llvm::StringRef HeaderPath instead.

clang-tools-extra/clangd/index/CanonicalIncludes.h

40 ↗

(On Diff #532964)

This time sounds a little late, since we don't want to rebuild the system header mapping every time we request a mapping for a particular header.

System header mapping is a singleton, we build it statically once on first request throughout the whole process and just use that for the rest of the execution.

clang-tools-extra/clangd/index/FileIndex.cpp

49 ↗

(On Diff #535740)

again you can have a const include_cleaner::PragmaIncludes & without shared_ptr here.

clang-tools-extra/clangd/index/FileIndex.h

118 ↗

(On Diff #535740)

you can have a const include_cleaner::PragmaIncludes& PI directly both in here and in indexHeaderSymbols.

after building a preamble we always have a non-null PragmaIncludes and these functions are always executed synchronously, so no need to extend ownership and make the code indirect.

clang-tools-extra/clangd/index/IndexAction.cpp

233

you can have a unique_ptr here instead (once we store a non-owning pointer in Opts.PragmaIncludes)

clang-tools-extra/clangd/index/SymbolCollector.cpp

250

yeah getLastRef is fine, i was mostly trying to avoid dodgy call to getOrCreateFileID.

you can also pass in just a llvm::StringRef HeaderPath now, as we're not using anything but file path. which you can retreive with getName from both FileEntry and FileEntryRef

412

s/SM.getFileEntryForID(FID)/FE

790

you can actually change include_cleaner::Macro to have a const IdentifierInfo* instead of the const_cast here. it's mostly an oversight to store a non-const IdentifierInfo in the first place.

837

once we have the optional<Header> as the value type you can also rewrite this as:

auto [It, Inserted] = SymbolProviders.try_emplace(S.ID);
if (Inserted) {
  auto Providers = include_cleaner::headersForSymbol(Sym, SM, Opts.PragmaIncludes.get());
  if(!Providers.empty()) It->second = Providers.front();
}

that way we can get rid of some redundant calls to headersForSymbol

839

std::move is unfortunately a little bit misleading at times. it is just a cast that returns an rvalue reference. hence on its own it doesn't actually move anything.

but it enables picking rvalue versions of constructors/assignment operators. (e.g. Foo(Foo&&) or operator=(Foo&&)) which can be implemented efficiently with the assumption that RHS object is being "moved from" and will no longer be used. that way certain types implement a move semantics by just consuming the RHS object. unfortunately that is not possbile when RHS is const, and expressions use regular copy semantics instead.

868

we're joining IncludeFiles and SymbolProviders, as they have the same keys.

right now you're iterating over the keys in IncludeFiles and doing a lookup into SymbolProviders using that key to get providers.
we can also do the iteration over SymbolProviders and do the lookup into IncludeFiles instead to find associated FID.

957

FWIW, i think we should assert(It != SymbolProviders.end()) (or IncludeFiles.end() when you change the outer loop). as we're inserting into these maps simultaneously.

960

we're not really getting any caching benefits as this map is inside the loop. can you put it outside the loop body?

963

we shouldn't continue here, it just means we can directly use SpellingIt to figure out what to put into NewSym.IncludeHeaders.

maybe something like:

auto [SpellingIt, Inserted] = HeaderSpelling.try_emplace(H);
if (Inserted)
  SpellingIt-> second = getSpelling(H, *PP);
if(!SpellingIt->second.empty()) {
  Symbol NewSym = *S;
  NewSym.IncludeHeaders.push_back({SpellingIt->second, 1, Directives});
  Symbols.insert(NewSym);
}

std::string getSpelling(const include_cleaner::Header &H, const Preprocessor &PP) {
  if(H.kind() == Physical) {
   // Check if we have a mapping for the header in system includes.
   // FIXME: get rid of this once stdlib mapping covers these system headers too.
   CanonicalIncludes CI;
   CI.addSystemHeadersMapping(PP.getLangOpts());
    if (auto Canonical = CI.mapHeader(H.physical()->getLastRef()); !Canonical.empty())
        return Canonical;
    if (!tooling::isSelfContainedHeader(...))
        return "";
  }
  // Otherwise use default spelling strategies in include_cleaner.
  return include_cleaner::spellHeader(H);
}

972

you can move the copy into this branch to make sure we're only doing that when necessary, e.g:

if(auto IncludeHeader = ...) {
  Symbol NewSym = *S;
  NewSym.IncludeHeaders...;
  Symbols.insert(NewSym);
}

clang-tools-extra/clangd/index/SymbolCollector.h

65

no need to have ownership here, we can have just a const include_cleaner::PragmaIncludes* PI. SymbolCollector shouldn't need to extend lifetimes

129

as mentioned above addSystemHeadersMapping is cheap after the first call in the process, we initialize the mappings only once. so you can just have it as a member in HeaderFileURICache, and drop it completely from the public interface of SymbolCollector

183

can we have a llvm::DenseMap<SymbolID, std::optional<include_cleaner::Header>>, to indicate there are no providers for a symbol, rather than missing the symbol completely from the mapping?

[VitaNuo]

Thank you!

clang-tools-extra/clangd/index/SymbolCollector.cpp

837

Sure, although I'm not sure where the redundant calls would be coming from. I've been under the impression that this function is supposed to be called once for each symbol.

868

Ok, makes sense now, thanks.

963

As agreed offline, it is not easy to split out a spelling function, since HeaderFileURIs is a private member of the symbol collector, and we seem to still need to generate URIs for physical headers going forward.

[VitaNuo]

Address review comments.

[kadircet]

thanks a lot for bearing with me, let's ship it!

clang-tools-extra/clangd/IncludeCleaner.cpp

91 ↗

(On Diff #539908)

nit: just const include_cleaner::PragmaIncludes* PI

clang-tools-extra/clangd/index/SymbolCollector.cpp

255

SysHeaderMapping doesn't need to be part of class state, just define it here completely. (also gets rid of the need for extra explanation and any possible questions around what would happen when we add system headers multiple times)

421

nit: can you move this above the call to mapCanonical and use it there too?

827–828

nit: early exits might be more readable, as we've more nesting now. (if (!Opts.Collect.. || shouldCollect..() == Symbol::Invalid) return;)

837

yeah unfortunately addDeclaration can be called multiple times (with increasing declaration quality, e.g. once for the forward declaration, then for a complete definition of a class)

949–952

nit: auto FID = IncludeFiles.at(SID);

949–952

if (!SpellingIt->second.empty()) {
  auto NewSym = *S;
  NewSym.IncludeHeaders.push_back({SpellingIt->second, 1, Directives});
  Symbols.insert(NewSym);
}

970

i believe this should be just H.physical()->getLastRef() ? as we're passing in a fileentryref into toURI

971

nit: braces, as we've multiple lines (despite being a single statement)

[VitaNuo]

Address the latest comment batch.

clang-tools-extra/clangd/index/IndexAction.cpp

IndexAction(std::shared_ptr<SymbolCollector> C,
               const index::IndexingOptions &Opts,
               std::function<void(SymbolSlab)> SymbolsCallback,
               std::function<void(RefSlab)> RefsCallback,
               std::function<void(RelationSlab)> RelationsCallback,
               std::function<void(IncludeGraph)> IncludeGraphCallback)
       : 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)), 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
       // reference are some examples). Since we index all the decls nested
       // inside, it becomes quadratic. So we give up on nested symbols.
       if (isDeeplyNested(D))
         return false;
       auto &SM = D->getASTContext().getSourceManager();
       auto FID = SM.getFileID(SM.getExpansionLoc(D->getLocation()));
       if (!FID.isValid())
         return true;
       return Collector->shouldIndexFile(FID);
     };
   }
 
   std::unique_ptr<ASTConsumer>
   CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
-    CI.getPreprocessor().addCommentHandler(PragmaHandler.get());
     Includes->addSystemHeadersMapping(CI.getLangOpts());
     if (IncludeGraphCallback != nullptr)
       CI.getPreprocessor().addPPCallbacks(
           std::make_unique<IncludeGraphCollector>(CI.getSourceManager(), IG));
 
     return index::createIndexingASTConsumer(Collector, Opts,
                                             CI.getPreprocessorPtr());
   }

 private:
   std::function<void(SymbolSlab)> SymbolsCallback;
   std::function<void(RefSlab)> RefsCallback;
   std::function<void(RelationSlab)> RelationsCallback;
   std::function<void(IncludeGraph)> IncludeGraphCallback;
   std::shared_ptr<SymbolCollector> Collector;
   std::unique_ptr<CanonicalIncludes> Includes;
   index::IndexingOptions Opts;
-  std::unique_ptr<CommentHandler> PragmaHandler;
   IncludeGraph IG;
 };
 
 } // namespace
 
 std::unique_ptr<FrontendAction> createStaticIndexingAction(
     SymbolCollector::Options Opts,
     std::function<void(SymbolSlab)> SymbolsCallback,
if (RefsCallback != nullptr) {
     Opts.RefFilter = RefKind::All;
     Opts.RefsInHeaders = true;
   }
   auto Includes = std::make_unique<CanonicalIncludes>();
   Opts.Includes = Includes.get();
   return std::make_unique<IndexAction>(
       std::make_shared<SymbolCollector>(std::move(Opts)), std::move(Includes),
       IndexOpts, SymbolsCallback, RefsCallback, RelationsCallback,
       IncludeGraphCallback);
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/SymbolCollector.h

 //===--- SymbolCollector.h ---------------------------------------*- C++-*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOLCOLLECTOR_H
 #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"
 #include "index/Symbol.h"
 #include "index/SymbolID.h"
 #include "index/SymbolOrigin.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 #include <functional>
 #include <optional>
 
 namespace clang {
 namespace clangd {
 
 /// Collect declarations (symbols) from an AST.
 /// It collects most declarations except:
struct Options {
     /// VFS that does not have absolute path), combine the fallback directory
     /// with symbols' paths to get absolute paths. This must be an absolute
     /// path.
     std::string FallbackDir;
     bool CollectIncludePath = false;
     /// If set, this is used to map symbol #include path to a potentially
     /// different #include path.
     const CanonicalIncludes *Includes = nullptr;
     // Populate the Symbol.References field.
     bool CountReferences = false;
     /// The symbol ref kinds that will be collected.
     /// If not set, SymbolCollector will not collect refs.
     /// Note that references of namespace decls are not collected, as they
     /// contribute large part of the index, and they are less useful compared
     /// with other decls.
     RefKind RefFilter = RefKind::Unknown;
     /// If set to true, SymbolCollector will collect all refs (from main file
public:
   // like call hierarchy.
   static const Decl *getRefContainer(const Decl *Enclosing,
                                      const SymbolCollector::Options &Opts);
 
   void initialize(ASTContext &Ctx) override;
 
   void setPreprocessor(std::shared_ptr<Preprocessor> PP) override {
     this->PP = PP.get();
+    PI.record(*this->PP);
+  }
+  void setPreprocessor(Preprocessor &PP) {
+    this->PP = &PP;
+    PI.record(*this->PP);
   }
-  void setPreprocessor(Preprocessor &PP) { this->PP = &PP; }
 
   bool
   handleDeclOccurrence(const Decl *D, index::SymbolRoleSet Roles,
                        ArrayRef<index::SymbolRelation> Relations,
                        SourceLocation Loc,
                        index::IndexDataConsumer::ASTNodeInfo ASTNode) override;
 
   bool handleMacroOccurrence(const IdentifierInfo *Name, const MacroInfo *MI,
private:
   std::optional<std::string> getIncludeHeader(const Symbol &S, FileID);
 
   SymbolID getSymbolIDCached(const Decl *D);
   SymbolID getSymbolIDCached(const llvm::StringRef MacroName,
                              const MacroInfo *MI, const SourceManager &SM);
 
   // 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 symbol ID to final include spelling.
+  // Needs to be persisted here because Symbols don't store any of their data.
+  llvm::DenseMap<SymbolID, std::string> SymbolIncludeSpelling;
   // 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:
   //   1) symbols declared in the preamble and referenced from the main file (
   //     which is not a header), or
   //   2) symbols declared and referenced from the main file (which is a header)
   RefSlab::Builder Refs;
   // All relations collected from the AST.
   RelationSlab::Builder Relations;
   ASTContext *ASTCtx;
   Preprocessor *PP = nullptr;
+  include_cleaner::PragmaIncludes PI;
   std::shared_ptr<GlobalCodeCompletionAllocator> CompletionAllocator;
   std::unique_ptr<CodeCompletionTUInfo> CompletionTUInfo;
   Options Opts;
   struct SymbolRef {
     SourceLocation Loc;
     FileID FID;
     index::SymbolRoleSet Roles;
     const Decl *Container;

clang-tools-extra/clangd/index/SymbolCollector.cpp

 //===--- SymbolCollector.cpp -------------------------------------*- C++-*-===//
 //
 // 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 "SymbolCollector.h"
 #include "AST.h"
 #include "CodeComplete.h"
 #include "CodeCompletionStrings.h"
 #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"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/DeclarationName.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Preprocessor.h"
 #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>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 /// If \p ND is a template specialization, returns the described template.
 /// Otherwise, returns \p ND.
 const NamedDecl &getTemplateOrThis(const NamedDecl &ND) {
class SymbolCollector::HeaderFileURICache {
   const CanonicalIncludes *Includes;
   llvm::StringRef FallbackDir;
   llvm::DenseMap<const FileEntry *, const std::string *> CacheFEToURI;
   llvm::StringMap<std::string> CachePathToURI;
   llvm::DenseMap<FileID, llvm::StringRef> CacheFIDToInclude;
   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) {
-  }
+                     const SymbolCollector::Options &Opts,
+                     const include_cleaner::PragmaIncludes &PI)
+      : PP(PP), SM(SM), Includes(Opts.Includes), FallbackDir(Opts.FallbackDir),
+        PI(PI) {}
 
   // Returns a canonical URI for the file \p FE.
   // We attempt to make the path absolute first.
   const std::string &toURI(const FileEntryRef FE) {
     auto R = CacheFEToURI.try_emplace(FE);
     if (R.second) {
       auto CanonPath = getCanonicalPath(FE, SM.getFileManager());
       R.first->second = &toURIInternal(CanonPath ? *CanonPath : FE.getName());
llvm::StringRef getIncludeHeader(FileID FID) {
     auto R = CacheFIDToInclude.try_emplace(FID);
     if (R.second)
       R.first->second = getIncludeHeaderUncached(FID);
     return R.first->second;
   }
 
 private:
   // This takes care of making paths absolute and path->URI caching, but no
   // FileManager-based canonicalization.
   const std::string &toURIInternal(llvm::StringRef Path) {
     auto R = CachePathToURI.try_emplace(Path);
     if (R.second) {
       llvm::SmallString<256> AbsPath = Path;
       if (!llvm::sys::path::is_absolute(AbsPath) && !FallbackDir.empty())
         llvm::sys::fs::make_absolute(FallbackDir, AbsPath);
       assert(llvm::sys::path::is_absolute(AbsPath) &&
              "If the VFS can't make paths absolute, a FallbackDir must be "
              "provided");
       llvm::sys::path::remove_dots(AbsPath, /*remove_dot_dot=*/true);
       R.first->second = URI::create(AbsPath).toString();
     }
     return R.first->second;
   }
 
   struct FrameworkHeaderPath {
std::optional<llvm::StringRef> getFrameworkHeaderIncludeSpelling(
     return llvm::StringRef(*CachedHeaderSpelling);
   }
 
   llvm::StringRef getIncludeHeaderUncached(FileID FID) {
     const auto FE = SM.getFileEntryRefForID(FID);
     if (!FE || FE->getName().empty())
       return "";
     llvm::StringRef Filename = FE->getName();
+
+    // 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;
+
     // 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);
       }
     }
     // Framework headers are spelled as <FrameworkName/Foo.h>, not
     // "path/FrameworkName.framework/Headers/Foo.h".
     auto &HS = PP->getHeaderSearchInfo();
     if (const auto *HFI = HS.getExistingFileInfo(*FE, /*WantExternal*/ false))
       if (!HFI->Framework.empty())
         if (auto Spelling =
                 getFrameworkHeaderIncludeSpelling(*FE, HFI->Framework, HS))
           return *Spelling;
 
     if (!tooling::isSelfContainedHeader(*FE, PP->getSourceManager(),
                                         PP->getHeaderSearchInfo())) {
       // A .inc or .def file is often included into a real header to define
       // symbols (e.g. LLVM tablegen files).
       if (Filename.endswith(".inc") || Filename.endswith(".def"))
         // Don't use cache reentrantly due to iterator invalidation.
         return getIncludeHeaderUncached(SM.getFileID(SM.getIncludeLoc(FID)));
       // Conservatively refuse to insert #includes to files without guards.
       return "";

 }
 
 SymbolCollector::SymbolCollector(Options Opts) : Opts(std::move(Opts)) {}
 SymbolCollector::~SymbolCollector() = default;
 
 void SymbolCollector::initialize(ASTContext &Ctx) {
   ASTCtx = &Ctx;
   HeaderFileURIs = std::make_unique<HeaderFileURICache>(
-      this->PP, ASTCtx->getSourceManager(), Opts);
+      this->PP, ASTCtx->getSourceManager(), Opts, PI);
   CompletionAllocator = std::make_shared<GlobalCodeCompletionAllocator>();
   CompletionTUInfo =
       std::make_unique<CodeCompletionTUInfo>(CompletionAllocator);
 }
 
 bool SymbolCollector::shouldCollectSymbol(const NamedDecl &ND,
                                           const ASTContext &ASTCtx,
                                           const Options &Opts,
bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
   std::string SnippetSuffix;
   getSignature(*CCS, &Signature, &SnippetSuffix, SymbolCompletion.Kind,
                SymbolCompletion.CursorKind);
   S.Signature = Signature;
   S.CompletionSnippetSuffix = SnippetSuffix;
 
   IndexedMacros.insert(Name);
   setIncludeLocation(S, DefLoc);
+  llvm::SmallVector<include_cleaner::Header> Headers =
+      include_cleaner::headersForSymbol(
+          include_cleaner::Macro{const_cast<IdentifierInfo *>(Name), DefLoc},
+          ASTCtx->getSourceManager(), &PI);
+  setSymbolProviders(S, Headers);
   Symbols.insert(S);
   return true;
 }
 
 void SymbolCollector::processRelations(
     const NamedDecl &ND, const SymbolID &ID,
     ArrayRef<index::SymbolRelation> Relations) {
   for (const auto &R : Relations) {
for (const auto &R : Relations) {
     if (*RKind == RelationKind::BaseOf)
       this->Relations.insert({ID, *RKind, ObjectID});
     else if (*RKind == RelationKind::OverriddenBy)
       this->Relations.insert({ObjectID, *RKind, ID});
   }
 }
 
 void SymbolCollector::setIncludeLocation(const Symbol &S, SourceLocation Loc) {
   if (Opts.CollectIncludePath &&
       shouldCollectIncludePath(S.SymInfo.Kind) != Symbol::Invalid)
     // Use the expansion location to get the #include header since this is
     // where the symbol is exposed.
     IncludeFiles[S.ID] =
         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] = Headers;
+}
+
 void SymbolCollector::finish() {
   // At the end of the TU, add 1 to the refcount of all referenced symbols.
   for (const auto &ID : ReferencedSymbols) {
     if (const auto *S = Symbols.find(ID)) {
       // SymbolSlab::Builder returns const symbols because strings are interned
       // and modifying returned symbols without inserting again wouldn't go
       // well. const_cast is safe here as we're modifying a data owned by the
       // Symbol. This reduces time spent in SymbolCollector by ~1%.
for (const IdentifierInfo *II : IndexedMacros) {
                 getSymbolIDCached(II->getName(), MI, PP->getSourceManager()))
           if (MI->isUsedForHeaderGuard())
             Symbols.erase(ID);
     }
   }
   llvm::DenseMap<FileID, bool> FileToContainsImportsOrObjC;
   // Fill in IncludeHeaders.
   // 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>";
         }
       }
-      // 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;
+      // Determine if the FID is #include'd or #import'ed.
+
+      Symbol NewSym;
+        NewSym = *S;
+        if (!IncludeHeader.empty()) {
           NewSym.IncludeHeaders.push_back({IncludeHeader, 1, Directives});
           Symbols.insert(NewSym);
-        }
+          continue;
+        }
+
+      // 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.
+        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) {
+            switch (H.kind()) {
+            case include_cleaner::Header::Kind::Verbatim:
+              SymbolIncludeSpelling[SID] = H.verbatim().str();
+              break;
+            case include_cleaner::Header::Kind::Standard:
+              SymbolIncludeSpelling[SID] = H.standard().name().str();
+              break;
+            case include_cleaner::Header::Kind::Physical:
+              SymbolIncludeSpelling[SID] = HeaderFileURIs->getIncludeHeader(
+                  ASTCtx->getSourceManager().getOrCreateFileID(
+                      H.physical(), SrcMgr::CharacteristicKind::C_User));
+              break;
+            }
+            if (!SymbolIncludeSpelling[SID].empty()) {
+              NewSym.IncludeHeaders.push_back(
+                  {SymbolIncludeSpelling[SID], 1, Directives});
+              break;
+            }
+          }
+        }
+        Symbols.insert(NewSym);
       }
     }
   }
 
   ReferencedSymbols.clear();
   IncludeFiles.clear();
+  SymbolProviders.clear();
   FilesWithObjCConstructs.clear();
 }
 
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
                                               bool IsMainFileOnly) {
   auto &Ctx = ND.getASTContext();
   auto &SM = Ctx.getSourceManager();
 
   Symbol S;
   S.ID = std::move(ID);
   std::string QName = printQualifiedName(ND);
   // FIXME: this returns foo:bar: for objective-C methods, we prefer only foo:
   // for consistency with CodeCompletionString and a clean name/signature split.
   std::tie(S.Scope, S.Name) = splitQualifiedName(QName);
   std::string TemplateSpecializationArgs = printTemplateSpecializationArgs(ND);
   S.TemplateSpecializationArgs = TemplateSpecializationArgs;
 
   // We collect main-file symbols, but do not use them for code completion.
   if (!IsMainFileOnly && isIndexedForCodeCompletion(ND, Ctx))
     S.Flags |= Symbol::IndexedForCodeCompletion;
   if (isImplementationDetail(&ND))
     S.Flags |= Symbol::ImplementationDetail;
   if (!IsMainFileOnly)
     S.Flags |= Symbol::VisibleOutsideFile;
   S.SymInfo = index::getSymbolInfo(&ND);
   auto Loc = nameLocation(ND, SM);
   assert(Loc.isValid() && "Invalid source location for NamedDecl");
const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
   if (S.Flags & Symbol::IndexedForCodeCompletion) {
     TypeStorage = OpaqueType::fromCompletionResult(*ASTCtx, SymbolCompletion);
     if (TypeStorage)
       S.Type = TypeStorage->raw();
   }
 
   Symbols.insert(S);
   setIncludeLocation(S, ND.getLocation());
+  llvm::SmallVector<include_cleaner::Header> Headers =
+      include_cleaner::headersForSymbol(include_cleaner::Symbol{ND},
+                                        ASTCtx->getSourceManager(), &PI);
+  setSymbolProviders(S, Headers);
   if (S.SymInfo.Lang == index::SymbolLanguage::ObjC)
     FilesWithObjCConstructs.insert(FID);
   return Symbols.find(S.ID);
 }
 
 void SymbolCollector::addDefinition(const NamedDecl &ND,
                                     const Symbol &DeclSym) {
   if (DeclSym.Definition)

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

 //===-- SymbolCollectorTests.cpp  -------------------------------*- C++ -*-===//
 //
 // 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 "Annotations.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "URI.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/FileSystemOptions.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
build(R"(
     }
   )");
   EXPECT_TRUE(shouldCollect("nx::Top_Level"));
   EXPECT_TRUE(shouldCollect("nx::Kind_Fine"));
 }
 
 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)
-          : DataConsumer(std::move(DataConsumer)), Opts(Opts),
-            PragmaHandler(PragmaHandler) {}
+                  const index::IndexingOptions &Opts)
+          : DataConsumer(std::move(DataConsumer)), Opts(Opts) {}
 
       std::unique_ptr<ASTConsumer>
       CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
-        if (PragmaHandler)
-          CI.getPreprocessor().addCommentHandler(PragmaHandler);
         return createIndexingASTConsumer(DataConsumer, Opts,
                                          CI.getPreprocessorPtr());
       }
 
       bool BeginInvocation(CompilerInstance &CI) override {
         // Make the compiler parse all comments.
         CI.getLangOpts().CommentOpts.ParseAllComments = true;
         return true;
       }
 
     private:
       std::shared_ptr<index::IndexDataConsumer> DataConsumer;
       index::IndexingOptions Opts;
-      CommentHandler *PragmaHandler;
     };
     index::IndexingOptions IndexOpts;
     IndexOpts.SystemSymbolFilter =
         index::IndexingOptions::SystemSymbolFilterKind::All;
     IndexOpts.IndexFunctionLocals = true;
     Collector = std::make_shared<SymbolCollector>(COpts);
-    return std::make_unique<IndexAction>(Collector, std::move(IndexOpts),
-                                         PragmaHandler);
+    return std::make_unique<IndexAction>(Collector, std::move(IndexOpts));
   }
 
   std::shared_ptr<SymbolCollector> Collector;
   SymbolCollector::Options COpts;
-  CommentHandler *PragmaHandler;
 };
 
 class SymbolCollectorTest : public ::testing::Test {
 public:
   SymbolCollectorTest()
       : InMemoryFileSystem(new llvm::vfs::InMemoryFileSystem),
         TestHeaderName(testPath("symbol.h")),
         TestFileName(testPath("symbol.cc")) {
     TestHeaderURI = URI::create(TestHeaderName).toString();
     TestFileURI = URI::create(TestFileName).toString();
   }
 
   // Note that unlike TestTU, no automatic header guard is added.
   // HeaderCode should start with #pragma once to be treated as modular.
   bool runSymbolCollector(llvm::StringRef HeaderCode, llvm::StringRef MainCode,
                           const std::vector<std::string> &ExtraArgs = {}) {
     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};
     Args.insert(Args.end(), ExtraArgs.begin(), ExtraArgs.end());
     // This allows to override the "-xc++" with something else, i.e.
     // -xobjective-c++.
     Args.push_back(TestFileName);
 
protected:
   std::string TestHeaderName;
   std::string TestHeaderURI;
   std::string TestFileName;
   std::string TestFileURI;
   SymbolSlab Symbols;
   RefSlab Refs;
   RelationSlab Relations;
   SymbolCollector::Options CollectorOpts;
-  std::unique_ptr<CommentHandler> PragmaHandler;
 };
 
 TEST_F(SymbolCollectorTest, CollectSymbols) {
   const std::string Header = R"(
     class Foo {
       Foo() {}
       Foo(int a) {}
       void f();
EXPECT_THAT(
                 includeHeader("<string>")),
           // Parameter names are demangled.
           AllOf(labeled("move(T &&value)"), includeHeader("<utility>")),
           AllOf(labeled("move(I, I, O)"), includeHeader("<algorithm>"))));
 }
 
 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 {};
   )";
   runSymbolCollector(Header, /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAre(
                            AllOf(qName("Foo"), declURI(TestHeaderURI),
                                  includeHeader("\"the/good/header.h\""))));
 }
 
 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 {};
   )";
   runSymbolCollector(Header, /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAre(
                            AllOf(qName("Foo"), declURI(TestHeaderURI),
                                  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();
   InMemoryFileSystem->addFile(IncFile, 0,
                               llvm::MemoryBuffer::getMemBuffer("class X {};"));
   llvm::IntrusiveRefCntPtr<FileManager> Files(
       new FileManager(FileSystemOptions(), InMemoryFileSystem));
   std::string HeaderCode = "#include \"test.inc\"\nclass Y {};";