1 files changed, 13 insertions, 44 deletions

diff --git a/clang-tools-extra/clangd/Preamble.cpp b/clang-tools-extra/clangd/Preamble.cpp
index d68c1ec97fa..41e3c1368eb 100644
--- a/clang-tools-extra/clangd/Preamble.cpp
+++ b/clang-tools-extra/clangd/Preamble.cpp
@@ -24,49 +24,14 @@ bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
          llvm::makeArrayRef(LHS.CommandLine).equals(RHS.CommandLine);
 }
 
-// This collects macro definitions in the *preamble region* of the main file.
-// (Contrast with CollectMainFileMacroExpansions in ParsedAST.cpp, which
-// collects macro *expansions* in the rest of the main file.
-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) {
-    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) {
-  }
+      : File(File), ParsedCallback(ParsedCallback) {}
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
-  std::vector<std::string> takeMainFileMacros() {
-    return std::move(MainFileMacros);
-  }
+  MainFileMacros takeMacros() { return std::move(Macros); }
 
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
@@ -79,14 +44,17 @@ public:
 
   void BeforeExecute(CompilerInstance &CI) override {
     CanonIncludes.addSystemHeadersMapping(CI.getLangOpts());
+    LangOpts = &CI.getLangOpts();
     SourceMgr = &CI.getSourceManager();
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
-    assert(SourceMgr && "SourceMgr must be set at this point");
+    assert(SourceMgr && LangOpts &&
+           "SourceMgr and LangOpts must be set at this point");
+
     return std::make_unique<PPChainedCallbacks>(
         collectIncludeStructureCallback(*SourceMgr, &Includes),
-        std::make_unique<CollectMainFileMacros>(*SourceMgr, &MainFileMacros));
+        std::make_unique<CollectMainFileMacros>(*SourceMgr, *LangOpts, Macros));
   }
 
   CommentHandler *getCommentHandler() override {
@@ -99,20 +67,21 @@ private:
   PreambleParsedCallback ParsedCallback;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
-  std::vector<std::string> MainFileMacros;
+  MainFileMacros Macros;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
-  SourceManager *SourceMgr = nullptr;
+  const clang::LangOptions *LangOpts = nullptr;
+  const SourceManager *SourceMgr = nullptr;
 };
 
 } // namespace
 
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<Diag> Diags, IncludeStructure Includes,
-                           std::vector<std::string> MainFileMacros,
+                           MainFileMacros Macros,
                            std::unique_ptr<PreambleFileStatusCache> StatCache,
                            CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Diags(std::move(Diags)),
-      Includes(std::move(Includes)), MainFileMacros(std::move(MainFileMacros)),
+      Includes(std::move(Includes)), Macros(std::move(Macros)),
       StatCache(std::move(StatCache)), CanonIncludes(std::move(CanonIncludes)) {
 }
 
@@ -181,7 +150,7 @@ buildPreamble(PathRef FileName, CompilerInvocation &CI,
     return std::make_shared<PreambleData>(
         std::move(*BuiltPreamble), std::move(Diags),
         SerializedDeclsCollector.takeIncludes(),
-        SerializedDeclsCollector.takeMainFileMacros(), std::move(StatCache),
+        SerializedDeclsCollector.takeMacros(), std::move(StatCache),
         SerializedDeclsCollector.takeCanonicalIncludes());
   } else {
     elog("Could not build a preamble for file {0}", FileName);