[clangd] Add helper for collecting #include directives in file.

Summary: Separate unit tests for the new function will be added in followup patch which will further refactor Headers.h

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: klimek, ilya-biryukov, MaskRay, jkorous, cfe-commits

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

llvm-svn: 332237

5 files changed, 81 insertions, 80 deletions

diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index fd8f1ec2765d..ab87463774c4 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -83,44 +83,13 @@ private:
   std::vector<const Decl *> TopLevelDecls;
 };
 
-class InclusionLocationsCollector : public PPCallbacks {
-public:
-  InclusionLocationsCollector(SourceManager &SourceMgr,
-                              InclusionLocations &IncLocations)
-      : SourceMgr(SourceMgr), IncLocations(IncLocations) {}
-
-  void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
-                          StringRef FileName, bool IsAngled,
-                          CharSourceRange FilenameRange, const FileEntry *File,
-                          StringRef SearchPath, StringRef RelativePath,
-                          const Module *Imported,
-                          SrcMgr::CharacteristicKind FileType) override {
-    auto SR = FilenameRange.getAsRange();
-    if (SR.isInvalid() || !File || File->tryGetRealPathName().empty())
-      return;
-
-    if (SourceMgr.isInMainFile(SR.getBegin())) {
-      // Only inclusion directives in the main file make sense. The user cannot
-      // select directives not in the main file.
-      IncLocations.emplace_back(halfOpenToRange(SourceMgr, FilenameRange),
-                                File->tryGetRealPathName());
-    }
-  }
-
-private:
-  SourceManager &SourceMgr;
-  InclusionLocations &IncLocations;
-};
-
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   std::vector<serialization::DeclID> takeTopLevelDeclIDs() {
     return std::move(TopLevelDeclIDs);
   }
 
-  InclusionLocations takeInclusionLocations() {
-    return std::move(IncLocations);
-  }
+  std::vector<Inclusion> takeInclusions() { return std::move(Inclusions); }
 
   void AfterPCHEmitted(ASTWriter &Writer) override {
     TopLevelDeclIDs.reserve(TopLevelDecls.size());
@@ -146,14 +115,15 @@ public:
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
     assert(SourceMgr && "SourceMgr must be set at this point");
-    return llvm::make_unique<InclusionLocationsCollector>(*SourceMgr,
-                                                          IncLocations);
+    return collectInclusionsInMainFileCallback(
+        *SourceMgr,
+        [this](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); });
   }
 
 private:
   std::vector<Decl *> TopLevelDecls;
   std::vector<serialization::DeclID> TopLevelDeclIDs;
-  InclusionLocations IncLocations;
+  std::vector<Inclusion> Inclusions;
   SourceManager *SourceMgr = nullptr;
 };
 
@@ -191,15 +161,15 @@ ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
     return llvm::None;
   }
 
-  InclusionLocations IncLocations;
+  std::vector<Inclusion> Inclusions;
   // Copy over the includes from the preamble, then combine with the
   // non-preamble includes below.
   if (Preamble)
-    IncLocations = Preamble->IncLocations;
+    Inclusions = Preamble->Inclusions;
 
-  Clang->getPreprocessor().addPPCallbacks(
-      llvm::make_unique<InclusionLocationsCollector>(Clang->getSourceManager(),
-                                                     IncLocations));
+  Clang->getPreprocessor().addPPCallbacks(collectInclusionsInMainFileCallback(
+      Clang->getSourceManager(),
+      [&Inclusions](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); }));
 
   if (!Action->Execute())
     log("Execute() failed when building AST for " + MainInput.getFile());
@@ -213,7 +183,7 @@ ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
   std::vector<const Decl *> ParsedDecls = Action->takeTopLevelDecls();
   return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
                    std::move(ParsedDecls), ASTDiags.take(),
-                   std::move(IncLocations));
+                   std::move(Inclusions));
 }
 
 void ParsedAST::ensurePreambleDeclsDeserialized() {
@@ -279,27 +249,27 @@ std::size_t ParsedAST::getUsedBytes() const {
          ::getUsedBytes(TopLevelDecls) + ::getUsedBytes(Diags);
 }
 
-const InclusionLocations &ParsedAST::getInclusionLocations() const {
-  return IncLocations;
+const std::vector<Inclusion> &ParsedAST::getInclusions() const {
+  return Inclusions;
 }
 
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<serialization::DeclID> TopLevelDeclIDs,
                            std::vector<Diag> Diags,
-                           InclusionLocations IncLocations)
+                           std::vector<Inclusion> Inclusions)
     : Preamble(std::move(Preamble)),
       TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)),
-      IncLocations(std::move(IncLocations)) {}
+      Inclusions(std::move(Inclusions)) {}
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      std::vector<const Decl *> TopLevelDecls,
-                     std::vector<Diag> Diags, InclusionLocations IncLocations)
+                     std::vector<Diag> Diags, std::vector<Inclusion> Inclusions)
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Diags(std::move(Diags)),
       TopLevelDecls(std::move(TopLevelDecls)), PreambleDeclsDeserialized(false),
-      IncLocations(std::move(IncLocations)) {
+      Inclusions(std::move(Inclusions)) {
   assert(this->Clang);
   assert(this->Action);
 }
@@ -448,8 +418,7 @@ CppFile::rebuildPreamble(CompilerInvocation &CI,
     return std::make_shared<PreambleData>(
         std::move(*BuiltPreamble),
         SerializedDeclsCollector.takeTopLevelDeclIDs(),
-        PreambleDiagnostics.take(),
-        SerializedDeclsCollector.takeInclusionLocations());
+        PreambleDiagnostics.take(), SerializedDeclsCollector.takeInclusions());
   } else {
     log("Could not build a preamble for file " + Twine(FileName));
     return nullptr;
diff --git a/clang-tools-extra/clangd/ClangdUnit.h b/clang-tools-extra/clangd/ClangdUnit.h
index 84b5acca76e2..64a19c0e61a9 100644
--- a/clang-tools-extra/clangd/ClangdUnit.h
+++ b/clang-tools-extra/clangd/ClangdUnit.h
@@ -12,6 +12,7 @@
 
 #include "Diagnostics.h"
 #include "Function.h"
+#include "Headers.h"
 #include "Path.h"
 #include "Protocol.h"
 #include "clang/Frontend/FrontendAction.h"
@@ -40,18 +41,18 @@ struct CompileCommand;
 
 namespace clangd {
 
-using InclusionLocations = std::vector<std::pair<Range, Path>>;
-
 // Stores Preamble and associated data.
 struct PreambleData {
   PreambleData(PrecompiledPreamble Preamble,
                std::vector<serialization::DeclID> TopLevelDeclIDs,
-               std::vector<Diag> Diags, InclusionLocations IncLocations);
+               std::vector<Diag> Diags, std::vector<Inclusion> Inclusions);
 
   PrecompiledPreamble Preamble;
   std::vector<serialization::DeclID> TopLevelDeclIDs;
   std::vector<Diag> Diags;
-  InclusionLocations IncLocations;
+  // Processes like code completions and go-to-definitions will need #include
+  // information, and their compile action skips preamble range.
+  std::vector<Inclusion> Inclusions;
 };
 
 /// Information required to run clang, e.g. to parse AST or do code completion.
@@ -95,14 +96,14 @@ public:
   /// Returns the esitmated size of the AST and the accessory structures, in
   /// bytes. Does not include the size of the preamble.
   std::size_t getUsedBytes() const;
-  const InclusionLocations &getInclusionLocations() const;
+  const std::vector<Inclusion> &getInclusions() const;
 
 private:
   ParsedAST(std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action,
             std::vector<const Decl *> TopLevelDecls, std::vector<Diag> Diags,
-            InclusionLocations IncLocations);
+            std::vector<Inclusion> Inclusions);
 
 private:
   void ensurePreambleDeclsDeserialized();
@@ -122,7 +123,7 @@ private:
   std::vector<Diag> Diags;
   std::vector<const Decl *> TopLevelDecls;
   bool PreambleDeclsDeserialized;
-  InclusionLocations IncLocations;
+  std::vector<Inclusion> Inclusions;
 };
 
 using ASTParsedCallback = std::function<void(PathRef Path, ParsedAST *)>;
diff --git a/clang-tools-extra/clangd/Headers.cpp b/clang-tools-extra/clangd/Headers.cpp
index af2445531795..58a81dcb2919 100644
--- a/clang-tools-extra/clangd/Headers.cpp
+++ b/clang-tools-extra/clangd/Headers.cpp
@@ -10,6 +10,7 @@
 #include "Headers.h"
 #include "Compiler.h"
 #include "Logger.h"
+#include "SourceCode.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
@@ -24,27 +25,34 @@ namespace {
 
 class RecordHeaders : public PPCallbacks {
 public:
-  RecordHeaders(llvm::StringSet<> &WrittenHeaders,
-                llvm::StringSet<> &ResolvedHeaders)
-      : WrittenHeaders(WrittenHeaders), ResolvedHeaders(ResolvedHeaders) {}
+  RecordHeaders(const SourceManager &SM,
+                std::function<void(Inclusion)> Callback)
+      : SM(SM), Callback(std::move(Callback)) {}
 
-  void InclusionDirective(SourceLocation /*HashLoc*/,
-                          const Token & /*IncludeTok*/,
+  // Record existing #includes - both written and resolved paths. Only #includes
+  // in the main file are collected.
+  void InclusionDirective(SourceLocation HashLoc, const Token & /*IncludeTok*/,
                           llvm::StringRef FileName, bool IsAngled,
-                          CharSourceRange /*FilenameRange*/,
-                          const FileEntry *File, llvm::StringRef /*SearchPath*/,
+                          CharSourceRange FilenameRange, const FileEntry *File,
+                          llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const Module * /*Imported*/,
                           SrcMgr::CharacteristicKind /*FileType*/) override {
-    WrittenHeaders.insert(
-        (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str());
-    if (File != nullptr && !File->tryGetRealPathName().empty())
-      ResolvedHeaders.insert(File->tryGetRealPathName());
+    // Only inclusion directives in the main file make sense. The user cannot
+    // select directives not in the main file.
+    if (HashLoc.isInvalid() || !SM.isInMainFile(HashLoc))
+      return;
+    std::string Written =
+        (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str();
+    std::string Resolved = (!File || File->tryGetRealPathName().empty())
+                               ? ""
+                               : File->tryGetRealPathName();
+    Callback({halfOpenToRange(SM, FilenameRange), Written, Resolved});
   }
 
 private:
-  llvm::StringSet<> &WrittenHeaders;
-  llvm::StringSet<> &ResolvedHeaders;
+  const SourceManager &SM;
+  std::function<void(Inclusion)> Callback;
 };
 
 } // namespace
@@ -58,6 +66,12 @@ bool HeaderFile::valid() const {
          (!Verbatim && llvm::sys::path::is_absolute(File));
 }
 
+std::unique_ptr<PPCallbacks>
+collectInclusionsInMainFileCallback(const SourceManager &SM,
+                                    std::function<void(Inclusion)> Callback) {
+  return llvm::make_unique<RecordHeaders>(SM, std::move(Callback));
+}
+
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.
 llvm::Expected<std::string>
@@ -110,17 +124,22 @@ calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
     return llvm::make_error<llvm::StringError>(
         "Failed to begin preprocessor only action for file " + File,
         llvm::inconvertibleErrorCode());
-  llvm::StringSet<> WrittenHeaders;
-  llvm::StringSet<> ResolvedHeaders;
-  Clang->getPreprocessor().addPPCallbacks(
-      llvm::make_unique<RecordHeaders>(WrittenHeaders, ResolvedHeaders));
+  std::vector<Inclusion> Inclusions;
+  Clang->getPreprocessor().addPPCallbacks(collectInclusionsInMainFileCallback(
+      Clang->getSourceManager(),
+      [&Inclusions](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); }));
   if (!Action.Execute())
     return llvm::make_error<llvm::StringError>(
         "Failed to execute preprocessor only action for file " + File,
         llvm::inconvertibleErrorCode());
+  llvm::StringSet<> IncludedHeaders;
+  for (const auto &Inc : Inclusions) {
+    IncludedHeaders.insert(Inc.Written);
+    if (!Inc.Resolved.empty())
+      IncludedHeaders.insert(Inc.Resolved);
+  }
   auto Included = [&](llvm::StringRef Header) {
-    return WrittenHeaders.find(Header) != WrittenHeaders.end() ||
-           ResolvedHeaders.find(Header) != ResolvedHeaders.end();
+    return IncludedHeaders.find(Header) != IncludedHeaders.end();
   };
   if (Included(DeclaringHeader.File) || Included(InsertedHeader.File))
     return "";
diff --git a/clang-tools-extra/clangd/Headers.h b/clang-tools-extra/clangd/Headers.h
index 3165dd5b1a9e..8aadc93a35a3 100644
--- a/clang-tools-extra/clangd/Headers.h
+++ b/clang-tools-extra/clangd/Headers.h
@@ -11,7 +11,9 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H
 
 #include "Path.h"
+#include "Protocol.h"
 #include "clang/Basic/VirtualFileSystem.h"
+#include "clang/Lex/PPCallbacks.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
@@ -32,6 +34,18 @@ struct HeaderFile {
   bool valid() const;
 };
 
+// An #include directive that we found in the main file.
+struct Inclusion {
+  Range R;             // Inclusion range.
+  std::string Written; // Inclusion name as written e.g. <vector>.
+  Path Resolved;       // Resolved path of included file. Empty if not resolved.
+};
+
+/// Returns a PPCallback that visits all inclusions in the main file.
+std::unique_ptr<PPCallbacks>
+collectInclusionsInMainFileCallback(const SourceManager &SM,
+                                    std::function<void(Inclusion)> Callback);
+
 /// Determines the preferred way to #include a file, taking into account the
 /// search path. Usually this will prefer a shorter representation like
 /// 'Foo/Bar.h' over a longer one like 'Baz/include/Foo/Bar.h'.
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index f70b6a6867c3..30f8ecf75bdf 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -234,12 +234,10 @@ std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
 
   std::vector<Location> Result;
   // Handle goto definition for #include.
-  for (auto &IncludeLoc : AST.getInclusionLocations()) {
-    Range R = IncludeLoc.first;
+  for (auto &Inc : AST.getInclusions()) {
     Position Pos = sourceLocToPosition(SourceMgr, SourceLocationBeg);
-
-    if (R.contains(Pos))
-      Result.push_back(Location{URIForFile{IncludeLoc.second}, {}});
+    if (!Inc.Resolved.empty() && Inc.R.contains(Pos))
+      Result.push_back(Location{URIForFile{Inc.Resolved}, {}});
   }
   if (!Result.empty())
     return Result;