[clangd] Fix header-guard check for include insertion, and don't index header guards.

[sammccall]

Summary

Both of these attempt to check whether a header guard exists while parsing the
file. However the file is only marked as guarded once clang finishes processing
it. We defer the checks and work until SymbolCollector::finish().

This is ugly and ad-hoc, deferring *all* work might be cleaner.

[kadircet]

Thanks!

mostly LGTM, only confused about changing symbolcollector to produce a unique header per symbol(where as it was(could) produce multiple headers before).
What is the rationale behind that one?

clangd/index/SymbolCollector.h

130

This is losing ability to store multiple header files. Is that intentional?

clangd/unittests/SymbolCollectorTests.cpp

1043

can you also add a case where we first see MACRO and then decl.

[sammccall]

Extract setIncludeLocation() and call it from a more obvious point in the code.
Tweak test to make it harder.

clangd/index/SymbolCollector.h

130

Careful reading of the code shows that ability never existed :-) addDeclaration always creates a new Symbol, sometimes populates its IncludeHeaders, and then replaces the existing symbol. We always find a single decl of the symbol we prefer in the TU (though sometimes it takes us a few attempts).

Of course, I never read the code that carefully - I wrote this as a vector<pair<SymbolID, FileID>> to "preserve" the ability to add multiple headers... and then the tests failed :-)

clangd/unittests/SymbolCollectorTests.cpp

1043

IIUC the idea is that macros are "eager-er" than decls so more likely to break the caching. That makes sense.

WDYT about just putting the macro first, so we *only* test the "hard" case.

[kadircet]

clangd/index/SymbolCollector.h

130

Ah, ok I see. Thanks for pointing this out!

clangd/unittests/SymbolCollectorTests.cpp

1043

exactly, SGTM

clangd/ClangdUnit.h

 } // namespace tooling
 
 namespace clangd {
 
 // Stores Preamble and associated data.
 struct PreambleData {
   PreambleData(PrecompiledPreamble Preamble, std::vector<Diag> Diags,
                IncludeStructure Includes,
+               std::vector<std::string> MainFileMacros,
                std::unique_ptr<PreambleFileStatusCache> StatCache,
                CanonicalIncludes CanonIncludes);
 
   tooling::CompileCommand CompileCommand;
   PrecompiledPreamble Preamble;
   std::vector<Diag> Diags;
   // Processes like code completions and go-to-definitions will need #include
   // information, and their compile action skips preamble range.
   IncludeStructure Includes;
+  // Macros defined in the preamble section of the main file.
+  // Users care about headers vs main-file, not preamble vs non-preamble.
+  // These should be treated as main-file entities e.g. for code completion.
+  std::vector<std::string> MainFileMacros;
   // Cache of FS operations performed when building the preamble.
   // When reusing a preamble, this cache can be consumed to save IO.
   std::unique_ptr<PreambleFileStatusCache> StatCache;
   CanonicalIncludes CanonIncludes;
 };
 
 /// Stores and provides access to parsed AST.
 class ParsedAST {

clangd/ClangdUnit.cpp

 #include "IncludeFixer.h"
 #include "Logger.h"
 #include "SourceCode.h"
 #include "Trace.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/LangOptions.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/MacroInfo.h"
protected:
   CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
     return llvm::make_unique<DeclTrackingASTConsumer>(/*ref*/ TopLevelDecls);
   }
 
 private:
   std::vector<Decl *> TopLevelDecls;
 };
 
+class CollectMainFileMacros : public PPCallbacks {
+public:
+  explicit CollectMainFileMacros(const SourceManager &SM,
+                                 std::vector<std::string> *Out)
+      : SM(SM), Out(Out) {}
+
+  void FileChanged(SourceLocation Loc, FileChangeReason,
+                   SrcMgr::CharacteristicKind, FileID Prev) {
+    InMainFile = SM.isWrittenInMainFile(Loc);
+  }
+
+  void MacroDefined(const Token &MacroName, const MacroDirective *MD) {
+    assert(MacroName.is(tok::identifier));
+    if (InMainFile)
+      MainFileMacros.insert(MacroName.getIdentifierInfo()->getName());
+  }
+
+  void EndOfMainFile() {
+    for (const auto& Entry : MainFileMacros)
+      Out->push_back(Entry.getKey());
+    llvm::sort(*Out);
+  }
+
+ private:
+  const SourceManager &SM;
+  bool InMainFile = true;
+  llvm::StringSet<> MainFileMacros;
+  std::vector<std::string> *Out;
+};
+
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
       : File(File), ParsedCallback(ParsedCallback) {
     addSystemHeadersMapping(&CanonIncludes);
   }
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
+  std::vector<std::string> takeMainFileMacros() {
+    return std::move(MainFileMacros);
+  }
+
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
   void AfterExecute(CompilerInstance &CI) override {
     if (!ParsedCallback)
       return;
     trace::Span Tracer("Running PreambleCallback");
     ParsedCallback(CI.getASTContext(), CI.getPreprocessorPtr(), CanonIncludes);
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     SourceMgr = &CI.getSourceManager();
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
     assert(SourceMgr && "SourceMgr must be set at this point");
-    return collectIncludeStructureCallback(*SourceMgr, &Includes);
+    return llvm::make_unique<PPChainedCallbacks>(
+        collectIncludeStructureCallback(*SourceMgr, &Includes),
+        llvm::make_unique<CollectMainFileMacros>(*SourceMgr, &MainFileMacros));
   }
 
   CommentHandler *getCommentHandler() override {
     IWYUHandler = collectIWYUHeaderMaps(&CanonIncludes);
     return IWYUHandler.get();
   }
 
 private:
   PathRef File;
   PreambleParsedCallback ParsedCallback;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
+  std::vector<std::string> MainFileMacros;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   SourceManager *SourceMgr = nullptr;
 };
 
 // When using a preamble, only preprocessor events outside its bounds are seen.
 // This is almost what we want: replaying transitive preprocessing wastes time.
 // However this confuses clang-tidy checks: they don't see any #includes!
 // So we replay the *non-transitive* #includes that appear in the main-file.

 }
 
 const CanonicalIncludes &ParsedAST::getCanonicalIncludes() const {
   return CanonIncludes;
 }
 
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<Diag> Diags, IncludeStructure Includes,
+                           std::vector<std::string> MainFileMacros,
                            std::unique_ptr<PreambleFileStatusCache> StatCache,
                            CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Diags(std::move(Diags)),
-      Includes(std::move(Includes)), StatCache(std::move(StatCache)),
-      CanonIncludes(std::move(CanonIncludes)) {}
+      Includes(std::move(Includes)), MainFileMacros(std::move(MainFileMacros)),
+      StatCache(std::move(StatCache)), CanonIncludes(std::move(CanonIncludes)) {
+}
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      std::vector<Decl *> LocalTopLevelDecls,
                      std::vector<Diag> Diags, IncludeStructure Includes,
                      CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
buildPreamble(PathRef FileName, CompilerInvocation &CI,
   CI.getFrontendOpts().SkipFunctionBodies = false;
 
   if (BuiltPreamble) {
     vlog("Built preamble of size {0} for file {1}", BuiltPreamble->getSize(),
          FileName);
     std::vector<Diag> Diags = PreambleDiagnostics.take();
     return std::make_shared<PreambleData>(
         std::move(*BuiltPreamble), std::move(Diags),
-        SerializedDeclsCollector.takeIncludes(), std::move(StatCache),
+        SerializedDeclsCollector.takeIncludes(),
+        SerializedDeclsCollector.takeMainFileMacros(), std::move(StatCache),
         SerializedDeclsCollector.takeCanonicalIncludes());
   } else {
     elog("Could not build a preamble for file {0}", FileName);
     return nullptr;
   }
 }
 
 llvm::Optional<ParsedAST>

clangd/CodeComplete.cpp

 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "clang/Lex/ExternalPreprocessorSource.h"
+#include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "clang/Sema/DeclSpec.h"
 #include "clang/Sema/Sema.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
struct SemaCompleteInput {
   PathRef FileName;
   const tooling::CompileCommand &Command;
   const PreambleData *Preamble;
   llvm::StringRef Contents;
   size_t Offset;
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS;
 };
 
+void loadMainFilePreambleMacros(const Preprocessor &PP,
+                                const PreambleData &Preamble) {
+  // The ExternalPreprocessorSource has our macros, if we know where to look.
+  // We can read all the macros using PreambleMacros->ReadDefinedMacros(),
+  // but this includes transitively included files, so may deserialize a lot.
+  ExternalPreprocessorSource *PreambleMacros = PP.getExternalSource();
+  // As we have the names of the macros, we can look up their IdentifierInfo
+  // and then use this to load just the macros we want.
+  IdentifierInfoLookup *PreambleIdentifiers =
+      PP.getIdentifierTable().getExternalIdentifierLookup();
+  if (!PreambleIdentifiers || !PreambleMacros)
+    return;
+  for (const auto& MacroName : Preamble.MainFileMacros)
+    if (auto *II = PreambleIdentifiers->get(MacroName))
+      PreambleMacros->updateOutOfDateIdentifier(*II);
+}
+
 // Invokes Sema code completion on a file.
 // If \p Includes is set, it will be updated based on the compiler invocation.
 bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
                       const SemaCompleteInput &Input,
                       IncludeStructure *Includes = nullptr) {
   trace::Span Tracer("Sema completion");
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS = Input.VFS;
bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   // The diagnostic options must be set before creating a CompilerInstance.
   CI->getDiagnosticOpts().IgnoreWarnings = true;
   // We reuse the preamble whether it's valid or not. This is a
   // correctness/performance tradeoff: building without a preamble is slow, and
   // completion is latency-sensitive.
   // However, if we're completing *inside* the preamble section of the draft,
   // overriding the preamble will break sema completion. Fortunately we can just
   // skip all includes in this case; these completions are really simple.
-  bool CompletingInPreamble =
-      ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0).Size >
-      Input.Offset;
+  PreambleBounds PreambleRegion =
+      ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
+  bool CompletingInPreamble = PreambleRegion.Size > Input.Offset;
   // NOTE: we must call BeginSourceFile after prepareCompilerInstance. Otherwise
   // the remapped buffers do not get freed.
   IgnoreDiagnostics DummyDiagsConsumer;
   auto Clang = prepareCompilerInstance(
       std::move(CI),
       (Input.Preamble && !CompletingInPreamble) ? &Input.Preamble->Preamble
                                                 : nullptr,
       std::move(ContentsBuffer), std::move(VFS), DummyDiagsConsumer);
   Clang->getPreprocessorOpts().SingleFileParseMode = CompletingInPreamble;
   Clang->setCodeCompletionConsumer(Consumer.release());
 
   SyntaxOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
     log("BeginSourceFile() failed when running codeComplete for {0}",
         Input.FileName);
     return false;
   }
+  // Macros can be defined within the preamble region of the main file.
+  // They don't fall nicely into our index/Sema dichotomy:
+  //  - they're not indexed for completion (they're not available across files)
+  //  - but Sema code complete won't see them: as part of the preamble, they're
+  //    deserialized only when mentioned.
+  // Force them to be deserialized so SemaCodeComplete sees them.
+  if (Input.Preamble)
+    loadMainFilePreambleMacros(Clang->getPreprocessor(), *Input.Preamble);
   if (Includes)
     Clang->getPreprocessor().addPPCallbacks(
         collectIncludeStructureCallback(Clang->getSourceManager(), Includes));
   if (!Action.Execute()) {
     log("Execute() failed when running codeComplete for {0}", Input.FileName);
     return false;
   }
   Action.EndSourceFile();

clangd/index/Symbol.h

public:
   class Builder {
   public:
     Builder() : UniqueStrings(Arena) {}
 
     /// Adds a symbol, overwriting any existing one with the same ID.
     /// This is a deep copy: underlying strings will be owned by the slab.
     void insert(const Symbol &S);
 
-    /// Returns the symbol with an ID, if it exists. Valid until next insert().
+    /// Removes the symbol with an ID, if it exists.
+    void erase(const SymbolID &ID) { Symbols.erase(ID); }
+
+    /// Returns the symbol with an ID, if it exists. Valid until insert/remove.
     const Symbol *find(const SymbolID &ID) {
-      auto I = SymbolIndex.find(ID);
-      return I == SymbolIndex.end() ? nullptr : &Symbols[I->second];
+      auto I = Symbols.find(ID);
+      return I == Symbols.end() ? nullptr : &I->second;
     }
 
     /// Consumes the builder to finalize the slab.
     SymbolSlab build() &&;
 
   private:
     llvm::BumpPtrAllocator Arena;
     /// Intern table for strings. Contents are on the arena.
     llvm::UniqueStringSaver UniqueStrings;
-    std::vector<Symbol> Symbols;
     /// Values are indices into Symbols vector.
-    llvm::DenseMap<SymbolID, size_t> SymbolIndex;
+    llvm::DenseMap<SymbolID, Symbol> Symbols;
   };
 
 private:
   SymbolSlab(llvm::BumpPtrAllocator Arena, std::vector<Symbol> Symbols)
       : Arena(std::move(Arena)), Symbols(std::move(Symbols)) {}
 
   llvm::BumpPtrAllocator Arena; // Owns Symbol data that the Symbols do not.
   std::vector<Symbol> Symbols;  // Sorted by SymbolID to allow lookup.
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_H

clangd/index/Symbol.cpp

 }
 
 // Copy the underlying data of the symbol into the owned arena.
 static void own(Symbol &S, llvm::UniqueStringSaver &Strings) {
   visitStrings(S, [&](llvm::StringRef &V) { V = Strings.save(V); });
 }
 
 void SymbolSlab::Builder::insert(const Symbol &S) {
-  auto R = SymbolIndex.try_emplace(S.ID, Symbols.size());
-  if (R.second) {
-    Symbols.push_back(S);
-    own(Symbols.back(), UniqueStrings);
-  } else {
-    auto &Copy = Symbols[R.first->second] = S;
-    own(Copy, UniqueStrings);
-  }
+  own(Symbols[S.ID] = S, UniqueStrings);
 }
 
 SymbolSlab SymbolSlab::Builder::build() && {
-  Symbols = {Symbols.begin(), Symbols.end()}; // Force shrink-to-fit.
-  // Sort symbols so the slab can binary search over them.
-  llvm::sort(Symbols,
+  // Sort symbols into vector so the slab can binary search over them.
+  std::vector<Symbol> SortedSymbols;
+  SortedSymbols.reserve(Symbols.size());
+  for (auto &Entry : Symbols)
+    SortedSymbols.push_back(std::move(Entry.second));
+  llvm::sort(SortedSymbols,
              [](const Symbol &L, const Symbol &R) { return L.ID < R.ID; });
   // We may have unused strings from overwritten symbols. Build a new arena.
   llvm::BumpPtrAllocator NewArena;
   llvm::UniqueStringSaver Strings(NewArena);
-  for (auto &S : Symbols)
+  for (auto &S : SortedSymbols)
     own(S, Strings);
-  return SymbolSlab(std::move(NewArena), std::move(Symbols));
+  return SymbolSlab(std::move(NewArena), std::move(SortedSymbols));
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/SymbolCollector.h

private:
 
   llvm::Optional<std::string> getIncludeHeader(llvm::StringRef QName, FileID);
   bool isSelfContainedHeader(FileID);
   // Heuristically headers that only want to be included via an umbrella.
   static bool isDontIncludeMeHeader(llvm::StringRef);
 
   // All Symbols collected from the AST.
   SymbolSlab::Builder Symbols;
+  // File IDs for Symbol.IncludeHeaders.
+  // The final spelling is calculated in finish().
+  llvm::DenseMap<SymbolID, FileID> IncludeFiles;
+  void setIncludeLocation(const Symbol &S, SourceLocation);
+  // Indexed macros, to be erased if they turned out to be include guards.
+  llvm::DenseSet<const IdentifierInfo *> IndexedMacros;
   // All refs collected from the AST.
   // Only symbols declared in preamble (from #include) and referenced from the
   // main file will be included.
   RefSlab::Builder Refs;
   ASTContext *ASTCtx;
   std::shared_ptr<Preprocessor> PP;
   std::shared_ptr<GlobalCodeCompletionAllocator> CompletionAllocator;
   std::unique_ptr<CodeCompletionTUInfo> CompletionTUInfo;

clangd/index/SymbolCollector.cpp

 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Index/IndexSymbol.h"
+#include "clang/Index/IndexingAction.h"
 #include "clang/Index/USRGeneration.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
bool SymbolCollector::handleMacroOccurence(const IdentifierInfo *Name,
                                            SourceLocation Loc) {
   if (!Opts.CollectMacro)
     return true;
   assert(PP.get());
 
   const auto &SM = PP->getSourceManager();
   auto DefLoc = MI->getDefinitionLoc();
 
-  // Header guards are not interesting in index. Builtin macros don't have
-  // useful locations and are not needed for code completions.
-  if (MI->isUsedForHeaderGuard() || MI->isBuiltinMacro())
+  // Builtin macros don't have useful locations and aren't needed in completion.
+  if (MI->isBuiltinMacro())
     return true;
 
   // Skip main-file symbols if we are not collecting them.
   bool IsMainFileSymbol = SM.isInMainFile(SM.getExpansionLoc(DefLoc));
   if (IsMainFileSymbol && !Opts.CollectMainFileSymbols)
     return false;
 
   // Also avoid storing predefined macros like __DBL_MIN__.
if (auto DeclLoc =
     S.CanonicalDeclaration = *DeclLoc;
 
   CodeCompletionResult SymbolCompletion(Name);
   const auto *CCS = SymbolCompletion.CreateCodeCompletionStringForMacro(
       *PP, *CompletionAllocator, *CompletionTUInfo);
   std::string Signature;
   std::string SnippetSuffix;
   getSignature(*CCS, &Signature, &SnippetSuffix);
-
-  std::string Include;
-  if (Opts.CollectIncludePath && shouldCollectIncludePath(S.SymInfo.Kind)) {
-    if (auto Header = getIncludeHeader(
-            Name->getName(), SM.getDecomposedExpansionLoc(DefLoc).first))
-      Include = std::move(*Header);
-  }
   S.Signature = Signature;
   S.CompletionSnippetSuffix = SnippetSuffix;
-  if (!Include.empty())
-    S.IncludeHeaders.emplace_back(Include, 1);
 
+  IndexedMacros.insert(Name);
+  setIncludeLocation(S, DefLoc);
   Symbols.insert(S);
   return true;
 }
 
+void SymbolCollector::setIncludeLocation(const Symbol &S,
+                                         SourceLocation Loc) {
+  if (Opts.CollectIncludePath)
+    if (shouldCollectIncludePath(S.SymInfo.Kind))
+      // 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::finish() {
   // At the end of the TU, add 1 to the refcount of all referenced symbols.
   auto IncRef = [this](const SymbolID &ID) {
     if (const auto *S = Symbols.find(ID)) {
       Symbol Inc = *S;
       ++Inc.References;
       Symbols.insert(Inc);
     }
   };
   for (const NamedDecl *ND : ReferencedDecls) {
     if (auto ID = getSymbolID(ND)) {
       IncRef(*ID);
     }
   }
   if (Opts.CollectMacro) {
     assert(PP);
+    // First, drop header guards. We can't identify these until EOF.
+    for (const IdentifierInfo *II : IndexedMacros) {
+      if (const auto *MI = PP->getMacroDefinition(II).getMacroInfo())
+        if (auto ID = getSymbolID(*II, MI, PP->getSourceManager()))
+          if (MI->isUsedForHeaderGuard())
+            Symbols.erase(*ID);
+    }
+    // Now increment refcounts.
     for (const IdentifierInfo *II : ReferencedMacros) {
       if (const auto *MI = PP->getMacroDefinition(II).getMacroInfo())
         if (auto ID = getSymbolID(*II, MI, PP->getSourceManager()))
           IncRef(*ID);
     }
   }
 
+  // Fill in IncludeHeaders.
+  // We delay this until end of TU so header guards are all resolved.
+  // Symbols in slabs aren' mutable, so insert() has to walk all the strings :-(
+  llvm::SmallString<256> QName;
+  for (const auto &Entry : IncludeFiles)
+    if (const Symbol *S = Symbols.find(Entry.first)) {
+      QName = S->Scope;
+      QName.append(S->Name);
+      if (auto Header = getIncludeHeader(QName, Entry.second)) {
+        Symbol NewSym = *S;
+        NewSym.IncludeHeaders.push_back({*Header, 1});
+        Symbols.insert(NewSym);
+      }
+    }
+
   const auto &SM = ASTCtx->getSourceManager();
   llvm::DenseMap<FileID, std::string> URICache;
   auto GetURI = [&](FileID FID) -> llvm::Optional<std::string> {
     auto Found = URICache.find(FID);
     if (Found == URICache.end()) {
       if (auto *FileEntry = SM.getFileEntryForID(FID)) {
         auto FileURI = toURI(SM, FileEntry->getName(), Opts);
         Found = URICache.insert({FID, FileURI}).first;
       } else {
         // Ignore cases where we can not find a corresponding file entry
         // for the loc, thoses are not interesting, e.g. symbols formed
         // via macro concatenation.
         return None;
       }
     }
     return Found->second;
   };
-
+  // Populate Refs slab from DeclRefs.
   if (auto MainFileURI = GetURI(SM.getMainFileID())) {
     for (const auto &It : DeclRefs) {
       if (auto ID = getSymbolID(It.first)) {
         for (const auto &LocAndRole : It.second) {
           auto FileID = SM.getFileID(LocAndRole.first);
           // FIXME: use the result to filter out references.
           shouldIndexFile(SM, FileID, Opts, &FilesToIndexCache);
           if (auto FileURI = GetURI(FileID)) {
if (auto MainFileURI = GetURI(SM.getMainFileID())) {
     }
   }
 
   ReferencedDecls.clear();
   ReferencedMacros.clear();
   DeclRefs.clear();
   FilesToIndexCache.clear();
   HeaderIsSelfContainedCache.clear();
+  IncludeFiles.clear();
 }
 
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
                                               bool IsMainFileOnly) {
   auto &Ctx = ND.getASTContext();
   auto &SM = Ctx.getSourceManager();
 
   Symbol S;
const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND, SymbolID ID,
   std::string Signature;
   std::string SnippetSuffix;
   getSignature(*CCS, &Signature, &SnippetSuffix);
   S.Signature = Signature;
   S.CompletionSnippetSuffix = SnippetSuffix;
   std::string ReturnType = getReturnType(*CCS);
   S.ReturnType = ReturnType;
 
-  std::string Include;
-  if (Opts.CollectIncludePath && shouldCollectIncludePath(S.SymInfo.Kind)) {
-    // Use the expansion location to get the #include header since this is
-    // where the symbol is exposed.
-    if (auto Header = getIncludeHeader(
-            QName, SM.getDecomposedExpansionLoc(ND.getLocation()).first))
-      Include = std::move(*Header);
-  }
-  if (!Include.empty())
-    S.IncludeHeaders.emplace_back(Include, 1);
-
   llvm::Optional<OpaqueType> TypeStorage;
   if (S.Flags & Symbol::IndexedForCodeCompletion) {
     TypeStorage = OpaqueType::fromCompletionResult(*ASTCtx, SymbolCompletion);
     if (TypeStorage)
       S.Type = TypeStorage->raw();
   }
 
   Symbols.insert(S);
+  setIncludeLocation(S, ND.getLocation());
   return Symbols.find(S.ID);
 }
 
 void SymbolCollector::addDefinition(const NamedDecl &ND,
                                     const Symbol &DeclSym) {
   if (DeclSym.Definition)
     return;
   // If we saw some forward declaration, we end up copying the symbol.

clangd/unittests/CodeCompleteTests.cpp

TEST(CompletionTest, MergeMacrosFromIndexAndSema) {
   Sym.ID = SymbolID("c:foo.cpp@8@macro@Clangd_Macro_Test");
   Sym.SymInfo.Kind = index::SymbolKind::Macro;
   Sym.Flags |= Symbol::IndexedForCodeCompletion;
   EXPECT_THAT(completions("#define Clangd_Macro_Test\nClangd_Macro_T^", {Sym})
                   .Completions,
               UnorderedElementsAre(Named("Clangd_Macro_Test")));
 }
 
-TEST(CompletionTest, NoMacroFromPreambleIfIndexIsSet) {
+TEST(CompletionTest, MacroFromPreamble) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  std::string FooHeader = testPath("foo.h");
+  FS.Files[FooHeader] = "#define CLANGD_PREAMBLE_HEADER x\n";
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer Server(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
   auto Results = completions(
-      R"cpp(#define CLANGD_PREAMBLE x
+      R"cpp(#include "foo.h"
+          #define CLANGD_PREAMBLE_MAIN x
 
           int x = 0;
           #define CLANGD_MAIN x
           void f() { CLANGD_^ }
       )cpp",
       {func("CLANGD_INDEX")});
-  // Index is overriden in code completion options, so the preamble symbol is
-  // not seen.
-  EXPECT_THAT(Results.Completions, UnorderedElementsAre(Named("CLANGD_MAIN"),
-                                                        Named("CLANGD_INDEX")));
+  // We should get results from the main file, including the preamble section.
+  // However no results from included files (the index should cover them).
+  EXPECT_THAT(Results.Completions,
+              UnorderedElementsAre(Named("CLANGD_PREAMBLE_MAIN"),
+                                   Named("CLANGD_MAIN"),
 }
 
 TEST(CompletionTest, DeprecatedResults) {
   std::string Body = R"cpp(
     void TestClangd();
     void TestClangc() __attribute__((deprecated("", "")));
   )cpp";
 

clangd/unittests/SymbolCollectorTests.cpp

 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
+#include "gmock/gmock-matchers.h"
 #include "gmock/gmock-more-matchers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include <memory>
 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using testing::_;
 using testing::AllOf;
 using testing::Contains;
+using testing::Each;
 using testing::ElementsAre;
 using testing::Field;
-using testing::IsEmpty;
 using testing::Not;
 using testing::Pair;
 using testing::UnorderedElementsAre;
 using testing::UnorderedElementsAreArray;
 
 // GMock helpers for matching Symbol.
 MATCHER_P(Labeled, Label, "") {
   return (arg.Name + arg.Signature).str() == Label;
TEST_F(SymbolCollectorTest, IncFileInNonHeader) {
   runSymbolCollector("", R"cpp(
     #include "test.inc"
   )cpp",
                      /*ExtraArgs=*/{"-I", testRoot()});
   EXPECT_THAT(Symbols, UnorderedElementsAre(AllOf(QName("X"), DeclURI(IncURI),
                                                   Not(IncludeHeader()))));
 }
 
+// Features that depend on header-guards are fragile. Header guards are only
+// recognized when the file ends, so we have to defer checking for them.
+TEST_F(SymbolCollectorTest, HeaderGuardDetected) {
+  CollectorOpts.CollectIncludePath = true;
+  CollectorOpts.CollectMacro = true;
+  runSymbolCollector(R"cpp(
+    #ifndef HEADER_GUARD_
+    #define HEADER_GUARD_
+
+    // Symbols are seen before the header guard is complete.
+    #define MACRO
+    int decl();
+
+    #endif // Header guard is recognized here.
+  )cpp",
+                     "");
+  EXPECT_THAT(Symbols, Not(Contains(QName("HEADER_GUARD_"))));
+  EXPECT_THAT(Symbols, Each(IncludeHeader()));
+}
+
 TEST_F(SymbolCollectorTest, NonModularHeader) {
   auto TU = TestTU::withHeaderCode("int x();");
   EXPECT_THAT(TU.headerSymbols(), ElementsAre(IncludeHeader()));
 
   // Files missing include guards aren't eligible for insertion.
   TU.ImplicitHeaderGuard = false;
   EXPECT_THAT(TU.headerSymbols(), ElementsAre(Not(IncludeHeader())));