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

[ioeric]

Summary

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

[ioeric]

• Rebase on origin/master

[ioeric]

• Unmerge D46670

[sammccall]

clangd/ClangdUnit.h

53

this might be a good place to explain *why* we collect these - processes like code completion will need #include information later, and their compile action skips the preamble range

clangd/Headers.h

47

as Ilya suggested, I like the idea of passing a callback(Inclusion) here rather than a vector reference. It's much more obvious at the callsite what's going to happen (and doesn't force the use of the vector).
Also lifetime issues are still a little complicated, but easier to debug I think.

[ioeric]

• Merged with origin/master
• Addressed review comments.

clangd/ClangdUnit.h

 //===--- ClangdUnit.h -------------------------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 
 #include "Diagnostics.h"
 #include "Function.h"
+#include "Headers.h"
 #include "Path.h"
 #include "Protocol.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Serialization/ASTBitCodes.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include <memory>

 }
 
 namespace tooling {
 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;
+  std::vector<Inclusion> Inclusions;
 };
 
 /// Information required to run clang, e.g. to parse AST or do code completion.
 struct ParseInputs {
   tooling::CompileCommand CompileCommand;
   IntrusiveRefCntPtr<vfs::FileSystem> FS;
   std::string Contents;
 };
public:
   /// this call might be expensive.
   ArrayRef<const Decl *> getTopLevelDecls();
 
   const std::vector<Diag> &getDiagnostics() const;
 
   /// 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();
 
   // In-memory preambles must outlive the AST, it is important that this member
   // goes before Clang and Action.
   std::shared_ptr<const PreambleData> Preamble;
   // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
   // on it) and CompilerInstance used to run it. That way we don't have to do
   // complex memory management of all Clang structures on our own. (They are
   // stored in CompilerInstance and cleaned up by
   // FrontendAction.EndSourceFile).
   std::unique_ptr<CompilerInstance> Clang;
   std::unique_ptr<FrontendAction> Action;
 
   // Data, stored after parsing.
   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 *)>;
 
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
 class CppFile {
 public:

clangd/ClangdUnit.cpp

std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override {
     return llvm::make_unique<DeclTrackingASTConsumer>(/*ref*/ TopLevelDecls);
   }
 
 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) 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());
     for (Decl *D : TopLevelDecls) {
       // Invalid top-level decls may not have been serialized.
       if (D->isInvalidDecl())
         continue;
       TopLevelDeclIDs.push_back(Writer.getDeclID(D));
public:
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     SourceMgr = &CI.getSourceManager();
   }
 
   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, Inclusions);
   }
 
 private:
   std::vector<Decl *> TopLevelDecls;
   std::vector<serialization::DeclID> TopLevelDeclIDs;
-  InclusionLocations IncLocations;
+  std::vector<Inclusion> Inclusions;
   SourceManager *SourceMgr = nullptr;
 };
 
 } // namespace
 
 void clangd::dumpAST(ParsedAST &AST, llvm::raw_ostream &OS) {
   AST.getASTContext().getTranslationUnitDecl()->dump(OS, true);
 }
ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
   auto Action = llvm::make_unique<ClangdFrontendAction>();
   const FrontendInputFile &MainInput = Clang->getFrontendOpts().Inputs[0];
   if (!Action->BeginSourceFile(*Clang, MainInput)) {
     log("BeginSourceFile() failed when building AST for " +
         MainInput.getFile());
     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));
 
   if (!Action->Execute())
     log("Execute() failed when building AST for " + MainInput.getFile());
 
   // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
   // has a longer lifetime.
   Clang->getDiagnostics().setClient(new IgnoreDiagnostics);
   // CompilerInstance won't run this callback, do it directly.
   ASTDiags.EndSourceFile();
 
   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() {
   if (PreambleDeclsDeserialized || !Preamble)
     return;
 
   std::vector<const Decl *> Resolved;
   Resolved.reserve(Preamble->TopLevelDeclIDs.size());

 std::size_t ParsedAST::getUsedBytes() const {
   auto &AST = getASTContext();
   // FIXME(ibiryukov): we do not account for the dynamically allocated part of
   // Message and Fixes inside each diagnostic.
   return AST.getASTAllocatedMemory() + AST.getSideTableAllocatedMemory() +
          ::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);
 }
 
 CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
                  std::shared_ptr<PCHContainerOperations> PCHs,
                  ASTParsedCallback ASTCallback)
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
CppFile::rebuildPreamble(CompilerInvocation &CI,
 
   if (BuiltPreamble) {
     log("Built preamble of size " + Twine(BuiltPreamble->getSize()) +
         " for file " + Twine(FileName));
 
     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;
   }
 }
 
 SourceLocation clangd::getBeginningOfIdentifier(ParsedAST &Unit,
                                                 const Position &Pos,

clangd/Headers.h

 //===--- Headers.h - Include headers -----------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_HEADERS_H
 #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"
 
 namespace clang {
 namespace clangd {
 
 /// Returns true if \p Include is literal include like "path" or <path>.
 bool isLiteralInclude(llvm::StringRef Include);
 
 /// Represents a header file to be #include'd.
 struct HeaderFile {
   std::string File;
   /// If this is true, `File` is a literal string quoted with <> or "" that
   /// can be #included directly; otherwise, `File` is an absolute file path.
   bool Verbatim;
 
   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 collects all inclusions in the main file.
+/// Inclusions are added to \p Includes.
+std::unique_ptr<PPCallbacks>
+collectInclusionsInMainFileCallback(const SourceManager &SM,
+                                    std::vector<Inclusion> &Inclusions);
+
 /// 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'.
 ///
 /// \param File is an absolute file path.
 /// \param DeclaringHeader is the original header corresponding to \p
 /// InsertedHeader e.g. the header that declares a symbol.
 /// \param InsertedHeader The preferred header to be inserted. This could be the

clangd/Headers.cpp

 //===--- Headers.cpp - Include headers ---------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 class RecordHeaders : public PPCallbacks {
 public:
-  RecordHeaders(llvm::StringSet<> &WrittenHeaders,
-                llvm::StringSet<> &ResolvedHeaders)
-      : WrittenHeaders(WrittenHeaders), ResolvedHeaders(ResolvedHeaders) {}
+  RecordHeaders(const SourceManager &SM, std::vector<Inclusion> &Includes)
+      : SM(SM), Includes(Includes) {}
 
-  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*/) 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();
+    Includes.push_back({halfOpenToRange(SM, FilenameRange), Written, Resolved});
   }
 
 private:
-  llvm::StringSet<> &WrittenHeaders;
-  llvm::StringSet<> &ResolvedHeaders;
+  const SourceManager &SM;
+  std::vector<Inclusion> &Includes;
 };
 
 } // namespace
 
 bool isLiteralInclude(llvm::StringRef Include) {
   return Include.startswith("<") || Include.startswith("\"");
 }
 
 bool HeaderFile::valid() const {
   return (Verbatim && isLiteralInclude(File)) ||
          (!Verbatim && llvm::sys::path::is_absolute(File));
 }
 
+std::unique_ptr<PPCallbacks>
+collectInclusionsInMainFileCallback(const SourceManager &SM,
+                                    std::vector<Inclusion> &Includes) {
+  return llvm::make_unique<RecordHeaders>(SM, Includes);
+}
+
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.
 llvm::Expected<std::string>
 calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
                      const HeaderFile &DeclaringHeader,
                      const HeaderFile &InsertedHeader,
                      const tooling::CompileCommand &CompileCommand,
                      IntrusiveRefCntPtr<vfs::FileSystem> FS) {
if (Clang->getFrontendOpts().Inputs.empty())
     return llvm::make_error<llvm::StringError>(
         "Empty frontend action inputs empty for file " + File,
         llvm::inconvertibleErrorCode());
   PreprocessOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     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));
   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 "";
 
   auto &HeaderSearchInfo = Clang->getPreprocessor().getHeaderSearchInfo();
   bool IsSystem = false;
 
   if (InsertedHeader.Verbatim)

clangd/XRefs.cpp

 std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
                                       const SymbolIndex *Index) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
   SourceLocation SourceLocationBeg =
       getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
 
   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;
 
   // Identified symbols at a specific position.
   auto Symbols = getSymbolAtPosition(AST, SourceLocationBeg);
 
   for (auto Item : Symbols.Macros) {