[clangd] Customizable URI schemes for dynamic index.

Summary:
This allows dynamic index to have consistent URI schemes with the
static index which can have customized URI schemes, which would make file
proximity scoring based on URIs easier.

Reviewers: sammccall

Reviewed By: sammccall

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

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

10 files changed, 89 insertions, 66 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index c2b68f8a970..3ca694c68e9 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -87,7 +87,8 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
                                    : getStandardResourceDir()),
-      FileIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
+      FileIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex(Opts.URISchemes)
+                                           : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index b6d5f160128..2da8518ea52 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -79,6 +79,10 @@ public:
     /// opened files and uses the index to augment code completion results.
     bool BuildDynamicSymbolIndex = false;
 
+    /// URI schemes to use when building the dynamic index.
+    /// If empty, the default schemes in SymbolCollector will be used.
+    std::vector<std::string> URISchemes;
+
     /// If set, use this index to augment code completion results.
     SymbolIndex *StaticIndex = nullptr;
 
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 7e19b00458f..ceb86b1b1c2 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -15,7 +15,8 @@
 namespace clang {
 namespace clangd {
 
-SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP) {
+SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                    llvm::ArrayRef<std::string> URISchemes) {
   SymbolCollector::Options CollectorOpts;
   // FIXME(ioeric): we might also want to collect include headers. We would need
   // to make sure all includes are canonicalized (with CanonicalIncludes), which
@@ -24,6 +25,8 @@ SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP) {
   // CommentHandler for IWYU pragma) to canonicalize includes.
   CollectorOpts.CollectIncludePath = false;
   CollectorOpts.CountReferences = false;
+  if (!URISchemes.empty())
+    CollectorOpts.URISchemes = URISchemes;
 
   SymbolCollector Collector(std::move(CollectorOpts));
   Collector.setPreprocessor(PP);
@@ -41,6 +44,9 @@ SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP) {
   return Collector.takeSymbols();
 }
 
+FileIndex::FileIndex(std::vector<std::string> URISchemes)
+    : URISchemes(std::move(URISchemes)) {}
+
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Slab) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Slab)
@@ -79,7 +85,7 @@ void FileIndex::update(PathRef Path, ASTContext *AST,
   } else {
     assert(PP);
     auto Slab = llvm::make_unique<SymbolSlab>();
-    *Slab = indexAST(*AST, PP);
+    *Slab = indexAST(*AST, PP, URISchemes);
     FSymbols.update(Path, std::move(Slab));
   }
   auto Symbols = FSymbols.allSymbols();
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index f25ed1749bc..1af37d19b5e 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -56,6 +56,10 @@ private:
 /// \brief This manages symbls from files and an in-memory index on all symbols.
 class FileIndex : public SymbolIndex {
 public:
+  /// If URISchemes is empty, the default schemes in SymbolCollector will be
+  /// used.
+  FileIndex(std::vector<std::string> URISchemes = {});
+
   /// \brief Update symbols in \p Path with symbols in \p AST. If \p AST is
   /// nullptr, this removes all symbols in the file.
   /// If \p AST is not null, \p PP cannot be null and it should be the
@@ -72,11 +76,14 @@ public:
 private:
   FileSymbols FSymbols;
   MemIndex Index;
+  std::vector<std::string> URISchemes;
 };
 
 /// Retrieves namespace and class level symbols in \p AST.
 /// Exposed to assist in unit tests.
-SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP);
+/// If URISchemes is empty, the default schemes in SymbolCollector will be used.
+SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                    llvm::ArrayRef<std::string> URISchemes = {});
 
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/unittests/clangd/ClangdTests.cpp b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
index ae47b60a53e..e98bc7fbcd9 100644
--- a/clang-tools-extra/unittests/clangd/ClangdTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
@@ -152,28 +152,6 @@ protected:
   }
 };
 
-constexpr const char* ClangdTestScheme = "ClangdTests";
-class TestURIScheme : public URIScheme {
-public:
-  llvm::Expected<std::string>
-  getAbsolutePath(llvm::StringRef /*Authority*/, llvm::StringRef Body,
-                  llvm::StringRef /*HintPath*/) const override {
-    llvm_unreachable("ClangdTests never makes absolute path.");
-  }
-
-  llvm::Expected<URI>
-  uriFromAbsolutePath(llvm::StringRef AbsolutePath) const override {
-    llvm_unreachable("ClangdTest never creates a test URI.");
-  }
-
-  llvm::Expected<std::string> getIncludeSpelling(const URI &U) const override {
-    return ("\"" + U.body().trim("/") + "\"").str();
-  }
-};
-
-static URISchemeRegistry::Add<TestURIScheme>
-    X(ClangdTestScheme, "Test scheme for ClangdTests.");
-
 TEST_F(ClangdVFSTest, Parse) {
   // FIXME: figure out a stable format for AST dumps, so that we can check the
   // output of the dump itself is equal to the expected one, not just that it's
diff --git a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
index 3d73ef3143d..1f00e4a4d30 100644
--- a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
@@ -96,6 +96,20 @@ void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   M.update(File.Filename, &AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
+TEST(FileIndexTest, CustomizedURIScheme) {
+  FileIndex M({"unittest"});
+  update(M, "f", "class string {};");
+
+  FuzzyFindRequest Req;
+  Req.Query = "";
+  bool SeenSymbol = false;
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+    EXPECT_EQ(Sym.CanonicalDeclaration.FileURI, "unittest:///f.h");
+    SeenSymbol = true;
+  });
+  EXPECT_TRUE(SeenSymbol);
+}
+
 TEST(FileIndexTest, IndexAST) {
   FileIndex M;
   update(M, "f1", "namespace ns { void f() {} class X {}; }");
diff --git a/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp b/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
index 19822b23a0b..531377dbc02 100644
--- a/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
+++ b/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
@@ -372,17 +372,15 @@ TEST_F(SymbolCollectorTest, SymbolRelativeWithFallback) {
               UnorderedElementsAre(AllOf(QName("Foo"), DeclURI(TestHeaderURI))));
 }
 
-#ifndef _WIN32
 TEST_F(SymbolCollectorTest, CustomURIScheme) {
   // Use test URI scheme from URITests.cpp
   CollectorOpts.URISchemes.insert(CollectorOpts.URISchemes.begin(), "unittest");
-  TestHeaderName = testPath("test-root/x.h");
-  TestFileName = testPath("test-root/x.cpp");
+  TestHeaderName = testPath("x.h");
+  TestFileName = testPath("x.cpp");
   runSymbolCollector("class Foo {};", /*Main=*/"");
-  EXPECT_THAT(Symbols,
-              UnorderedElementsAre(AllOf(QName("Foo"), DeclURI("unittest:x.h"))));
+  EXPECT_THAT(Symbols, UnorderedElementsAre(
+                           AllOf(QName("Foo"), DeclURI("unittest:///x.h"))));
 }
-#endif
 
 TEST_F(SymbolCollectorTest, InvalidURIScheme) {
   // Use test URI scheme from URITests.cpp
diff --git a/clang-tools-extra/unittests/clangd/TestFS.cpp b/clang-tools-extra/unittests/clangd/TestFS.cpp
index 9020c6da63b..000d55854dc 100644
--- a/clang-tools-extra/unittests/clangd/TestFS.cpp
+++ b/clang-tools-extra/unittests/clangd/TestFS.cpp
@@ -7,7 +7,11 @@
 //
 //===----------------------------------------------------------------------===//
 #include "TestFS.h"
+#include "URI.h"
+#include "clang/AST/DeclCXX.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
+#include "llvm/Support/Path.h"
 #include "gtest/gtest.h"
 
 namespace clang {
@@ -62,5 +66,38 @@ std::string testPath(PathRef File) {
   return Path.str();
 }
 
+/// unittest: is a scheme that refers to files relative to testRoot()
+class TestScheme : public URIScheme {
+public:
+  static const char *Scheme;
+
+  llvm::Expected<std::string>
+  getAbsolutePath(llvm::StringRef /*Authority*/, llvm::StringRef Body,
+                  llvm::StringRef HintPath) const override {
+    assert(HintPath.startswith(testRoot()));
+    llvm::SmallString<16> Path(Body.begin(), Body.end());
+    llvm::sys::path::native(Path);
+    return testPath(Path);
+  }
+
+  llvm::Expected<URI>
+  uriFromAbsolutePath(llvm::StringRef AbsolutePath) const override {
+    llvm::StringRef Body = AbsolutePath;
+    if (!Body.consume_front(testRoot()))
+      return llvm::make_error<llvm::StringError>(
+          AbsolutePath + "does not start with " + testRoot(),
+          llvm::inconvertibleErrorCode());
+
+    return URI(Scheme, /*Authority=*/"",
+               llvm::sys::path::convert_to_slash(Body));
+  }
+};
+
+const char *TestScheme::Scheme = "unittest";
+
+static URISchemeRegistry::Add<TestScheme> X(TestScheme::Scheme, "Test schema");
+
+volatile int UnittestSchemeAnchorSource = 0;
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/unittests/clangd/TestFS.h b/clang-tools-extra/unittests/clangd/TestFS.h
index 06a1d8aca3d..be4aac4f792 100644
--- a/clang-tools-extra/unittests/clangd/TestFS.h
+++ b/clang-tools-extra/unittests/clangd/TestFS.h
@@ -56,6 +56,11 @@ const char *testRoot();
 // Returns a suitable absolute path for this OS.
 std::string testPath(PathRef File);
 
+// unittest: is a scheme that refers to files relative to testRoot()
+// This anchor is used to force the linker to link in the generated object file
+// and thus register unittest: URI scheme plugin.
+extern volatile int UnittestSchemeAnchorSource;
+
 } // namespace clangd
 } // namespace clang
 #endif
diff --git a/clang-tools-extra/unittests/clangd/URITests.cpp b/clang-tools-extra/unittests/clangd/URITests.cpp
index 4ddc1645afb..f6ac0c6b7f0 100644
--- a/clang-tools-extra/unittests/clangd/URITests.cpp
+++ b/clang-tools-extra/unittests/clangd/URITests.cpp
@@ -14,6 +14,11 @@
 
 namespace clang {
 namespace clangd {
+
+// Force the unittest URI scheme to be linked,
+static int LLVM_ATTRIBUTE_UNUSED UnittestSchemeAnchorDest =
+    UnittestSchemeAnchorSource;
+
 namespace {
 
 using ::testing::AllOf;
@@ -22,38 +27,6 @@ MATCHER_P(Scheme, S, "") { return arg.scheme() == S; }
 MATCHER_P(Authority, A, "") { return arg.authority() == A; }
 MATCHER_P(Body, B, "") { return arg.body() == B; }
 
-// Assume all files in the schema have a "test-root/" root directory, and the
-// schema path is the relative path to the root directory.
-// So the schema of "/some-dir/test-root/x/y/z" is "test:x/y/z".
-class TestScheme : public URIScheme {
-public:
-  static const char *Scheme;
-
-  static const char *TestRoot;
-
-  llvm::Expected<std::string>
-  getAbsolutePath(llvm::StringRef /*Authority*/, llvm::StringRef Body,
-                  llvm::StringRef HintPath) const override {
-    auto Pos = HintPath.find(TestRoot);
-    assert(Pos != llvm::StringRef::npos);
-    return (HintPath.substr(0, Pos + llvm::StringRef(TestRoot).size()) + Body)
-        .str();
-  }
-
-  llvm::Expected<URI>
-  uriFromAbsolutePath(llvm::StringRef AbsolutePath) const override {
-    auto Pos = AbsolutePath.find(TestRoot);
-    assert(Pos != llvm::StringRef::npos);
-    return URI(Scheme, /*Authority=*/"",
-               AbsolutePath.substr(Pos + llvm::StringRef(TestRoot).size()));
-  }
-};
-
-const char *TestScheme::Scheme = "unittest";
-const char *TestScheme::TestRoot = "/test-root/";
-
-static URISchemeRegistry::Add<TestScheme> X(TestScheme::Scheme, "Test schema");
-
 std::string createOrDie(llvm::StringRef AbsolutePath,
                         llvm::StringRef Scheme = "file") {
   auto Uri = URI::create(AbsolutePath, Scheme);
@@ -167,12 +140,12 @@ TEST(URITest, Resolve) {
 #else
   EXPECT_EQ(resolveOrDie(parseOrDie("file:/a/b/c")), "/a/b/c");
   EXPECT_EQ(resolveOrDie(parseOrDie("file://auth/a/b/c")), "/a/b/c");
-  EXPECT_EQ(resolveOrDie(parseOrDie("unittest:a/b/c"), "/dir/test-root/x/y/z"),
-            "/dir/test-root/a/b/c");
   EXPECT_THAT(resolveOrDie(parseOrDie("file://au%3dth/%28x%29/y/%20z")),
               "/(x)/y/ z");
   EXPECT_THAT(resolveOrDie(parseOrDie("file:///c:/x/y/z")), "c:/x/y/z");
 #endif
+  EXPECT_EQ(resolveOrDie(parseOrDie("unittest:a"), testPath("x")),
+            testPath("a"));
 }
 
 TEST(URITest, Platform) {