[clangd] Customizable URI schemes for dynamic index.

[ioeric]

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.

[sammccall]

Not ignoring this patch: the implementation LG, whether it's the right design depends on where D47935 ends up I think.
(I suspect this is what we want, is there any urgency to landing this part?)

[sammccall]

unittests/clangd/FileIndexTests.cpp

99

This can use the shared unittest: scheme from the followup patch, right?

Add a fixme or lift that into this patch?

[ioeric]

• Pulled unittest URI refactoring from D47935

[ioeric]

unittests/clangd/FileIndexTests.cpp

99

Done. Lifted the TestFS change into this patch.

clangd/ClangdServer.h

struct Options {
 
     /// Cached preambles are potentially large. If false, store them on disk.
     bool StorePreamblesInMemory = true;
 
     /// If true, ClangdServer builds a dynamic in-memory index for symbols in
     /// 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;
 
     /// The resource directory is used to find internal headers, overriding
     /// defaults and -resource-dir compiler flag).
     /// If None, ClangdServer calls CompilerInvocation::GetResourcePath() to
     /// obtain the standard resource directory.
     llvm::Optional<StringRef> ResourceDir = llvm::None;

clangd/ClangdServer.cpp

 
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            FileSystemProvider &FSProvider,
                            DiagnosticsConsumer &DiagConsumer,
                            const Options &Opts)
     : CompileArgs(CDB, Opts.ResourceDir ? Opts.ResourceDir->str()
                                         : getStandardResourceDir()),
       DiagConsumer(DiagConsumer), FSProvider(FSProvider),
-      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
       // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(
           Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,

clangd/index/FileIndex.h

private:
 
   /// \brief Stores the latest snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSlabs;
 };
 
 /// \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
   /// preprocessor that was used to build \p AST.
   void update(PathRef Path, ASTContext *AST, std::shared_ptr<Preprocessor> PP);
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void lookup(const LookupRequest &Req,
               llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 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
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clangd/index/FileIndex.cpp

 #include "FileIndex.h"
 #include "SymbolCollector.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Preprocessor.h"
 
 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
   // is not trivial given the current way of collecting symbols: we only have
   // AST at this point, but we also need preprocessor callbacks (e.g.
   // 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);
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
   std::vector<const Decl *> TopLevelDecls(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
   index::indexTopLevelDecls(AST, TopLevelDecls, Collector, IndexOpts);
 
   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)
     FileToSlabs.erase(Path);
   else
     FileToSlabs[Path] = std::move(Slab);
 }
 

 
 void FileIndex::update(PathRef Path, ASTContext *AST,
                        std::shared_ptr<Preprocessor> PP) {
   if (!AST) {
     FSymbols.update(Path, nullptr);
   } 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();
   Index.build(std::move(Symbols));
 }
 
 bool FileIndex::fuzzyFind(
     const FuzzyFindRequest &Req,

unittests/clangd/ClangdTests.cpp

std::string parseSourceAndDumpAST(
     Server.addDocument(SourceFilename, SourceContents);
     auto Result = dumpASTWithoutMemoryLocs(Server, SourceFilename);
     EXPECT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for diagnostics";
     EXPECT_EQ(ExpectErrors, DiagConsumer.hadErrorInLastDiags());
     return Result;
   }
 };
 
-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
   // different.
   auto Empty = parseSourceAndDumpAST("foo.cpp", "", {});
   auto OneDecl = parseSourceAndDumpAST("foo.cpp", "int a;", {});
   auto SomeDecls = parseSourceAndDumpAST("foo.cpp", "int a; int b; int c;", {});
   EXPECT_NE(Empty, OneDecl);

unittests/clangd/FileIndexTests.cpp

void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   TestTU File;
   File.Filename = (Basename + ".cpp").str();
   File.HeaderFilename = (Basename + ".h").str();
   File.HeaderCode = Code;
   auto AST = File.build();
   M.update(File.Filename, &AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
+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 {
+    return URI::parse(HardcodedURI);
+  }
+  static const char HardcodedURI[];
+  static const char Scheme[];
+};
+
+const char TestURIScheme::HardcodedURI[] = "ClangdTests:///test";
+const char TestURIScheme::Scheme[] = "ClangdTest";
+
+static URISchemeRegistry::Add<TestURIScheme> X(TestURIScheme::Scheme,
+                                               "Test scheme for ClangdTests.");
+
+TEST(FileIndexTest, CustomizedURIScheme) {
+  FileIndex M({TestURIScheme::Scheme});
+  update(M, "f", "class string {};");
+
+  FuzzyFindRequest Req;
+  Req.Query = "";
+  bool SeenSymbol = false;
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+    EXPECT_EQ(Sym.CanonicalDeclaration.FileURI, TestURIScheme::HardcodedURI);
+    SeenSymbol = true;
+  });
+  EXPECT_TRUE(SeenSymbol);
+}
+
 TEST(FileIndexTest, IndexAST) {
   FileIndex M;
   update(M, "f1", "namespace ns { void f() {} class X {}; }");
 
   FuzzyFindRequest Req;
   Req.Query = "";
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(M, Req), UnorderedElementsAre("ns::f", "ns::X"));