[clangd] Fixes for #include insertion.

Summary:
o Avoid inserting a header include into the header itself.
o Avoid inserting non-header files (by not indexing symbols in main
files at all).
o Canonicalize include paths for symbols in dynamic index.

Reviewers: sammccall, ilya-biryukov

Reviewed By: ilya-biryukov

Subscribers: klimek, jkorous-apple, cfe-commits

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

llvm-svn: 325523

9 files changed, 94 insertions, 39 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 5cd310bf04c..ab2eea6b215 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -327,13 +327,11 @@ ClangdServer::insertInclude(PathRef File, StringRef Code,
     if (!Resolved)
       return Resolved.takeError();
 
-    auto FS = FSProvider.getTaggedFileSystem(File).Value;
     tooling::CompileCommand CompileCommand =
         CompileArgs.getCompileCommand(File);
-    FS->setCurrentWorkingDirectory(CompileCommand.Directory);
-
     auto Include =
-        shortenIncludePath(File, Code, *Resolved, CompileCommand, FS);
+        calculateIncludePath(File, Code, *Resolved, CompileCommand,
+                             FSProvider.getTaggedFileSystem(File).Value);
     if (!Include)
       return Include.takeError();
     if (Include->empty())
diff --git a/clang-tools-extra/clangd/Headers.cpp b/clang-tools-extra/clangd/Headers.cpp
index faf73c7f99f..db8029b2ac0 100644
--- a/clang-tools-extra/clangd/Headers.cpp
+++ b/clang-tools-extra/clangd/Headers.cpp
@@ -16,6 +16,7 @@
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
@@ -45,10 +46,17 @@ private:
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.
 llvm::Expected<std::string>
-shortenIncludePath(llvm::StringRef File, llvm::StringRef Code,
-                   llvm::StringRef Header,
-                   const tooling::CompileCommand &CompileCommand,
-                   IntrusiveRefCntPtr<vfs::FileSystem> FS) {
+calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
+                     llvm::StringRef Header,
+                     const tooling::CompileCommand &CompileCommand,
+                     IntrusiveRefCntPtr<vfs::FileSystem> FS) {
+  assert(llvm::sys::path::is_absolute(File) &&
+         llvm::sys::path::is_absolute(Header));
+
+  if (File == Header)
+    return "";
+  FS->setCurrentWorkingDirectory(CompileCommand.Directory);
+
   // Set up a CompilerInstance and create a preprocessor to collect existing
   // #include headers in \p Code. Preprocesor also provides HeaderSearch with
   // which we can calculate the shortest include path for \p Header.
diff --git a/clang-tools-extra/clangd/Headers.h b/clang-tools-extra/clangd/Headers.h
index 57fd523465f..5e74ed62ffe 100644
--- a/clang-tools-extra/clangd/Headers.h
+++ b/clang-tools-extra/clangd/Headers.h
@@ -18,18 +18,21 @@
 
 namespace clang {
 namespace clangd {
+
 /// 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 Header is an absolute file path.
-/// \return A quoted "path" or <path>. If \p Header is already (directly)
-/// included in the file (including those included via different paths), this
-/// returns an empty string.
+/// \return A quoted "path" or <path>. This returns an empty string if:
+///   - \p Header is already (directly) included in the file (including those
+///   included via different paths).
+///   - \p Header is the same as \p File.
 llvm::Expected<std::string>
-shortenIncludePath(PathRef File, llvm::StringRef Code, llvm::StringRef Header,
-                   const tooling::CompileCommand &CompileCommand,
-                   IntrusiveRefCntPtr<vfs::FileSystem> FS);
+calculateIncludePath(PathRef File, llvm::StringRef Code, llvm::StringRef Header,
+                     const tooling::CompileCommand &CompileCommand,
+                     IntrusiveRefCntPtr<vfs::FileSystem> FS);
 
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/CanonicalIncludes.cpp b/clang-tools-extra/clangd/index/CanonicalIncludes.cpp
index f805cfafb83..0ac858e3d56 100644
--- a/clang-tools-extra/clangd/index/CanonicalIncludes.cpp
+++ b/clang-tools-extra/clangd/index/CanonicalIncludes.cpp
@@ -28,6 +28,7 @@ void CanonicalIncludes::addRegexMapping(llvm::StringRef RE,
 }
 
 llvm::StringRef CanonicalIncludes::mapHeader(llvm::StringRef Header) const {
+  std::lock_guard<std::mutex> Lock(RegexMutex);
   for (auto &Entry : RegexHeaderMappingTable) {
 #ifndef NDEBUG
     std::string Dummy;
diff --git a/clang-tools-extra/clangd/index/CanonicalIncludes.h b/clang-tools-extra/clangd/index/CanonicalIncludes.h
index 1cfb10c435c..5f9c043885f 100644
--- a/clang-tools-extra/clangd/index/CanonicalIncludes.h
+++ b/clang-tools-extra/clangd/index/CanonicalIncludes.h
@@ -23,6 +23,7 @@
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Regex.h"
+#include <mutex>
 #include <string>
 #include <vector>
 
@@ -31,6 +32,7 @@ namespace clangd {
 
 /// Maps a definition location onto an #include file, based on a set of filename
 /// rules.
+/// Only const methods (i.e. mapHeader) in this class are thread safe.
 class CanonicalIncludes {
 public:
   CanonicalIncludes() = default;
@@ -53,6 +55,8 @@ private:
   // arbitrary regexes.
   mutable std::vector<std::pair<llvm::Regex, std::string>>
       RegexHeaderMappingTable;
+  // Guards Regex matching as it's not thread-safe.
+  mutable std::mutex RegexMutex;
 };
 
 /// Returns a CommentHandler that parses pragma comment on include files to
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 1bdd8093bb6..3f1cdeb2a12 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -15,14 +15,30 @@ namespace clang {
 namespace clangd {
 namespace {
 
+const CanonicalIncludes *canonicalIncludesForSystemHeaders() {
+  static const auto *Includes = [] {
+    auto *I = new CanonicalIncludes();
+    addSystemHeadersMapping(I);
+    return I;
+  }();
+  return Includes;
+}
+
 /// Retrieves namespace and class level symbols in \p Decls.
 std::unique_ptr<SymbolSlab> indexAST(ASTContext &Ctx,
                                      std::shared_ptr<Preprocessor> PP,
                                      llvm::ArrayRef<const Decl *> Decls) {
   SymbolCollector::Options CollectorOpts;
-  // Code completion gets main-file results from Sema.
-  // But we leave this option on because features like go-to-definition want it.
-  CollectorOpts.IndexMainFiles = true;
+  // Although we do not index symbols in main files (e.g. cpp file), information
+  // in main files like definition locations of class declarations will still be
+  // collected; thus, the index works for go-to-definition.
+  // FIXME(ioeric): handle IWYU pragma for dynamic index. We might want to make
+  // SymbolCollector always provides include canonicalization (e.g. IWYU, STL).
+  // FIXME(ioeric): get rid of `IndexMainFiles` as this is always set to false.
+  CollectorOpts.IndexMainFiles = false;
+  CollectorOpts.CollectIncludePath = true;
+  CollectorOpts.Includes = canonicalIncludesForSystemHeaders();
+
   auto Collector = std::make_shared<SymbolCollector>(std::move(CollectorOpts));
   Collector->setPreprocessor(std::move(PP));
   index::IndexingOptions IndexOpts;
diff --git a/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp b/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
index 77ff77f6547..d4c9d7d1e29 100644
--- a/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
+++ b/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
@@ -580,8 +580,11 @@ TEST(CompletionTest, DynamicIndexMultiFile) {
                       /*StorePreamblesInMemory=*/true,
                       /*BuildDynamicSymbolIndex=*/true);
 
-  Server.addDocument(testPath("foo.cpp"), R"cpp(
+  FS.Files[testPath("foo.h")] = R"cpp(
       namespace ns { class XYZ {}; void foo(int x) {} }
+  )cpp";
+  Server.addDocument(testPath("foo.cpp"), R"cpp(
+      #include "foo.h"
   )cpp");
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for preamble";
 
diff --git a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
index 5016da40f54..04bdf9a5c44 100644
--- a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestFS.h"
 #include "index/FileIndex.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
@@ -83,17 +84,28 @@ std::vector<std::string> match(const SymbolIndex &I,
 }
 
 /// Create an ParsedAST for \p Code. Returns None if \p Code is empty.
-llvm::Optional<ParsedAST> build(std::string Path, llvm::StringRef Code) {
+/// \p Code is put into <Path>.h which is included by \p <BasePath>.cpp.
+llvm::Optional<ParsedAST> build(llvm::StringRef BasePath,
+                                llvm::StringRef Code) {
   if (Code.empty())
     return llvm::None;
-  const char *Args[] = {"clang", "-xc++", Path.c_str()};
+
+  assert(llvm::sys::path::extension(BasePath).empty() &&
+         "BasePath must be a base file path without extension.");
+  llvm::IntrusiveRefCntPtr<vfs::InMemoryFileSystem> VFS(
+      new vfs::InMemoryFileSystem);
+  std::string Path = (BasePath + ".cpp").str();
+  std::string Header = (BasePath + ".h").str();
+  VFS->addFile(Path, 0, llvm::MemoryBuffer::getMemBuffer(""));
+  VFS->addFile(Header, 0, llvm::MemoryBuffer::getMemBuffer(Code));
+  const char *Args[] = {"clang", "-xc++", "-include", Header.c_str(),
+                        Path.c_str()};
 
   auto CI = createInvocationFromCommandLine(Args);
 
   auto Buf = llvm::MemoryBuffer::getMemBuffer(Code);
   auto AST = ParsedAST::Build(std::move(CI), nullptr, std::move(Buf),
-                              std::make_shared<PCHContainerOperations>(),
-                              vfs::getRealFileSystem());
+                              std::make_shared<PCHContainerOperations>(), VFS);
   assert(AST.hasValue());
   return std::move(*AST);
 }
@@ -169,6 +181,20 @@ TEST(FileIndexTest, IgnoreClassMembers) {
   EXPECT_THAT(match(M, Req), UnorderedElementsAre("X"));
 }
 
+#ifndef LLVM_ON_WIN32
+TEST(FileIndexTest, CanonicalizeSystemHeader) {
+  FileIndex M;
+  std::string File = testPath("bits/basic_string");
+  M.update(File, build(File, "class string {};").getPointer());
+
+  FuzzyFindRequest Req;
+  Req.Query = "";
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+    EXPECT_EQ(Sym.Detail->IncludeHeader, "<string>");
+  });
+}
+#endif
+
 } // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/unittests/clangd/HeadersTests.cpp b/clang-tools-extra/unittests/clangd/HeadersTests.cpp
index bd4567ad277..2fad1d5ca9e 100644
--- a/clang-tools-extra/unittests/clangd/HeadersTests.cpp
+++ b/clang-tools-extra/unittests/clangd/HeadersTests.cpp
@@ -24,24 +24,14 @@ public:
   }
 
 protected:
-  // Adds \p File with content \p Code to the sub directory and returns the
-  // absolute path.
-//  std::string addToSubdir(PathRef File, llvm::StringRef Code = "") {
-//    assert(llvm::sys::path::is_relative(File) && "FileName should be relative");
-//    llvm::SmallString<32> Path;
-//    llvm::sys::path::append(Path, Subdir, File);
-//    FS.Files[Path] = Code;
-//    return Path.str();
-//  };
-
-  // Shortens the header and returns "" on error.
-  std::string shorten(PathRef Header) {
+  // Calculates the include path for \p Header, or returns "" on error.
+  std::string calculate(PathRef Header) {
     auto VFS = FS.getTaggedFileSystem(MainFile).Value;
     auto Cmd = CDB.getCompileCommand(MainFile);
     assert(static_cast<bool>(Cmd));
     VFS->setCurrentWorkingDirectory(Cmd->Directory);
     auto Path =
-        shortenIncludePath(MainFile, FS.Files[MainFile], Header, *Cmd, VFS);
+        calculateIncludePath(MainFile, FS.Files[MainFile], Header, *Cmd, VFS);
     if (!Path) {
       llvm::consumeError(Path.takeError());
       return std::string();
@@ -58,7 +48,7 @@ protected:
 TEST_F(HeadersTest, InsertInclude) {
   std::string Path = testPath("sub/bar.h");
   FS.Files[Path] = "";
-  EXPECT_EQ(shorten(Path), "\"bar.h\"");
+  EXPECT_EQ(calculate(Path), "\"bar.h\"");
 }
 
 TEST_F(HeadersTest, DontInsertDuplicateSameName) {
@@ -67,7 +57,7 @@ TEST_F(HeadersTest, DontInsertDuplicateSameName) {
 )cpp";
   std::string BarHeader = testPath("sub/bar.h");
   FS.Files[BarHeader] = "";
-  EXPECT_EQ(shorten(BarHeader), "");
+  EXPECT_EQ(calculate(BarHeader), "");
 }
 
 TEST_F(HeadersTest, DontInsertDuplicateDifferentName) {
@@ -76,7 +66,7 @@ TEST_F(HeadersTest, DontInsertDuplicateDifferentName) {
   FS.Files[MainFile] = R"cpp(
 #include "sub/bar.h"  // not shortest
 )cpp";
-  EXPECT_EQ(shorten(BarHeader), "");
+  EXPECT_EQ(calculate(BarHeader), "");
 }
 
 TEST_F(HeadersTest, StillInsertIfTrasitivelyIncluded) {
@@ -89,7 +79,13 @@ TEST_F(HeadersTest, StillInsertIfTrasitivelyIncluded) {
   FS.Files[MainFile] = R"cpp(
 #include "bar.h"
 )cpp";
-  EXPECT_EQ(shorten(BazHeader), "\"baz.h\"");
+  EXPECT_EQ(calculate(BazHeader), "\"baz.h\"");
+}
+
+TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
+  MainFile = testPath("main.h");
+  FS.Files[MainFile] = "";
+  EXPECT_EQ(calculate(MainFile), "");
 }
 
 } // namespace