[clangd] Make use of URIs in FileShardedIndex

[kadircet]

Summary

This makes FileShardedIndex more robust and gets rid of the need for a
URIToFileCache, as it is done by the callers now and it is only once per file,
rather than once per symbol.

[kadircet]

• Return None and change PathRef to auto since it is URI now.

[sammccall]

LG but I think we can skip one lot of conversion, let me know what you think.

clang-tools-extra/clangd/index/Background.cpp

182–183

is this really important enough to be worth saying twice :-)

clang-tools-extra/clangd/index/FileIndex.cpp

386

I don't actually think it's important that we use filenames as the PreambleSymbols sharding key. We can skip the whole resolving thing and just use the URIs. We should go through and update docs to change "path" to "key" in FileSymbols I guess. Am I missing something?

Sorry I think I alluded to this a bit cryptically...

[kadircet]

• Update FileSymbols to express keys are not necessarily file paths.
• Use Uri for keys of PreambleSymbols

[kadircet]

clang-tools-extra/clangd/index/Background.cpp

182–183

copy pasta is the worst kind of pasta :D

[sammccall]

clang-tools-extra/clangd/index/FileIndex.h

59

may want to mention that keys are typically files, up to you

clang-tools-extra/clangd/index/Background.cpp

 
 /// Given index results from a TU, only update symbols coming from files that
 /// are different or missing from than \p ShardVersionsSnapshot. Also stores new
 /// index information on IndexStorage.
 void BackgroundIndex::update(
     llvm::StringRef MainFile, IndexFileIn Index,
     const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
     bool HadErrors) {
-  llvm::StringMap<FileDigest> FilesToUpdate;
+  // Keys are URIs.
+  llvm::StringMap<std::pair<Path, FileDigest>> FilesToUpdate;
+  // Note that sources do not contain any information regarding missing headers,
+  // since we don't even know what absolute path they should fall in.
   for (const auto &IndexIt : *Index.Sources) {
     const auto &IGN = IndexIt.getValue();
-    // Note that sources do not contain any information regarding missing
-    // headers, since we don't even know what absolute path they should fall in.
-    auto AbsPath = llvm::cantFail(URI::resolve(IGN.URI, MainFile),
-                                  "Failed to resovle URI");
-    const auto DigestIt = ShardVersionsSnapshot.find(AbsPath);
+    auto AbsPath = URI::resolve(IGN.URI, MainFile);
+    if (!AbsPath) {
+      elog("Failed to resolve URI: {0}", AbsPath.takeError());
+      continue;
+    }
+    const auto DigestIt = ShardVersionsSnapshot.find(*AbsPath);
     // File has different contents, or indexing was successful this time.
     if (DigestIt == ShardVersionsSnapshot.end() ||
         DigestIt->getValue().Digest != IGN.Digest ||
         (DigestIt->getValue().HadErrors && !HadErrors))
-      FilesToUpdate[AbsPath] = IGN.Digest;
+      FilesToUpdate[IGN.URI] = {std::move(*AbsPath), IGN.Digest};
   }
 
   // Shard slabs into files.
-  FileShardedIndex ShardedIndex(std::move(Index), MainFile);
+  FileShardedIndex ShardedIndex(std::move(Index));
 
   // Build and store new slabs for each updated file.
   for (const auto &FileIt : FilesToUpdate) {
-    PathRef Path = FileIt.first();
-    auto IF = ShardedIndex.getShard(Path);
+    auto Uri = FileIt.first();
+    // ShardedIndex should always have a shard for a file in Index.Sources.
+    auto IF = ShardedIndex.getShard(Uri).getValue();
+    PathRef Path = FileIt.getValue().first;
 
     // Only store command line hash for main files of the TU, since our
     // current model keeps only one version of a header file.
     if (Path != MainFile)
       IF.Cmd.reset();
 
     // We need to store shards before updating the index, since the latter
     // consumes slabs.
     // FIXME: Also skip serializing the shard if it is already up-to-date.
     if (auto Error = IndexStorageFactory(Path)->storeShard(Path, IF))
       elog("Failed to write background-index shard for file {0}: {1}", Path,
            std::move(Error));
 
     {
       std::lock_guard<std::mutex> Lock(ShardVersionsMu);
-      const auto &Hash = FileIt.getValue();
+      const auto &Hash = FileIt.getValue().second;
       auto DigestIt = ShardVersions.try_emplace(Path);
       ShardVersion &SV = DigestIt.first->second;
       // Skip if file is already up to date, unless previous index was broken
       // and this one is not.
       if (!DigestIt.second && SV.Digest == Hash && SV.HadErrors && !HadErrors)
         continue;
       SV.Digest = Hash;
       SV.HadErrors = HadErrors;

clang-tools-extra/clangd/index/FileIndex.h

 /// How to handle duplicated symbols across multiple files.
 enum class DuplicateHandling {
   // Pick a random symbol. Less accurate but faster.
   PickOne,
   // Merge symbols. More accurate but slower.
   Merge,
 };
 
-/// A container of Symbols from several source files. It can be updated
-/// at source-file granularity, replacing all symbols from one file with a new
-/// set.
+/// A container of slabs associated with a key. It can be updated at key
+/// granularity, replacing all slabs belonging to a key with a new set. Keys are
+/// usually file paths/uris.
 ///
-/// This implements a snapshot semantics for symbols in a file. Each update to a
-/// file will create a new snapshot for all symbols in the file. Snapshots are
-/// managed with shared pointers that are shared between this class and the
-/// users. For each file, this class only stores a pointer pointing to the
-/// newest snapshot, and an outdated snapshot is deleted by the last owner of
-/// the snapshot, either this class or the symbol index.
+/// This implements snapshot semantics. Each update will create a new snapshot
+/// for all slabs of the Key. Snapshots are managed with shared pointers that
+/// are shared between this class and the users. For each key, this class only
+/// stores a pointer pointing to the newest snapshot, and an outdated snapshot
+/// is deleted by the last owner of the snapshot, either this class or the
+/// symbol index.
 ///
 /// The snapshot semantics keeps critical sections minimal since we only need
 /// locking when we swap or obtain references to snapshots.
 class FileSymbols {
 public:
-  /// Updates all symbols and refs in a file.
-  /// If either is nullptr, corresponding data for \p Path will be removed.
-  /// If CountReferences is true, \p Refs will be used for counting References
+  /// Updates all slabs associated with the \p Key.
+  /// If either is nullptr, corresponding data for \p Key will be removed.
+  /// If CountReferences is true, \p Refs will be used for counting references
   /// during merging.
-  void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
+  void update(llvm::StringRef Key, std::unique_ptr<SymbolSlab> Symbols,
               std::unique_ptr<RefSlab> Refs,
               std::unique_ptr<RelationSlab> Relations, bool CountReferences);
 
-  /// The index keeps the symbols alive.
+  /// The index keeps the slabs alive.
   /// Will count Symbol::References based on number of references in the main
   /// files, while building the index with DuplicateHandling::Merge option.
   std::unique_ptr<SymbolIndex>
   buildIndex(IndexType,
              DuplicateHandling DuplicateHandle = DuplicateHandling::PickOne);
 
 private:
   struct RefSlabAndCountReferences {
     std::shared_ptr<RefSlab> Slab;
     bool CountReferences = false;
   };
   mutable std::mutex Mutex;
 
-  /// Stores the latest symbol snapshots for all active files.
-  llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSymbols;
-  /// Stores the latest ref snapshots for all active files.
-  llvm::StringMap<RefSlabAndCountReferences> FileToRefs;
-  /// Stores the latest relation snapshots for all active files.
-  llvm::StringMap<std::shared_ptr<RelationSlab>> FileToRelations;
+  llvm::StringMap<std::shared_ptr<SymbolSlab>> SymbolsSnapshot;
+  llvm::StringMap<RefSlabAndCountReferences> RefsSnapshot;
+  llvm::StringMap<std::shared_ptr<RelationSlab>> RelatiosSnapshot;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
 public:
   FileIndex(bool UseDex = true);
 
SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
                              std::shared_ptr<Preprocessor> PP,
                              const CanonicalIncludes &Includes);
 
 /// Takes slabs coming from a TU (multiple files) and shards them per
 /// declaration location.
 struct FileShardedIndex {
   /// \p HintPath is used to convert file URIs stored in symbols into absolute
   /// paths.
-  explicit FileShardedIndex(IndexFileIn Input, PathRef HintPath);
+  explicit FileShardedIndex(IndexFileIn Input);
 
-  /// Returns absolute paths for all files that has a shard.
-  std::vector<PathRef> getAllFiles() const;
+  /// Returns uris for all files that has a shard.
+  std::vector<llvm::StringRef> getAllSources() const;
 
-  /// Generates index shard for the \p File. Note that this function results in
+  /// Generates index shard for the \p Uri. Note that this function results in
   /// a copy of all the relevant data.
   /// Returned index will always have Symbol/Refs/Relation Slabs set, even if
   /// they are empty.
-  IndexFileIn getShard(PathRef File) const;
+  llvm::Optional<IndexFileIn> getShard(llvm::StringRef Uri) const;
 
 private:
   // Contains all the information that belongs to a single file.
   struct FileShard {
     // Either declared or defined in the file.
     llvm::DenseSet<const Symbol *> Symbols;
     // Reference occurs in the file.
     llvm::DenseSet<const Ref *> Refs;
     // Subject is declared in the file.
     llvm::DenseSet<const Relation *> Relations;
     // Contains edges for only the direct includes.
     IncludeGraph IG;
   };
 
   // Keeps all the information alive.
   const IndexFileIn Index;
-  // Mapping from absolute paths to slab information.
+  // Mapping from URIs to slab information.
   llvm::StringMap<FileShard> Shards;
   // Used to build RefSlabs.
   llvm::DenseMap<const Ref *, SymbolID> RefToSymID;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clang-tools-extra/clangd/index/FileIndex.cpp

 #include "support/Logger.h"
 #include "support/Path.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include <memory>
 #include <tuple>
 #include <utility>
 #include <vector>
vlog("indexed {0} AST for {1} version {2}:\n"
        "  relations slab: {8} relations, {9} bytes",
        IsIndexMainAST ? "file" : "preamble", FileName, Version, Syms.size(),
        Syms.bytes(), Refs.size(), Refs.numRefs(), Refs.bytes(),
        Relations.size(), Relations.bytes());
   return std::make_tuple(std::move(Syms), std::move(Refs),
                          std::move(Relations));
 }
 
-// Resolves URI to file paths with cache.
-class URIToFileCache {
-public:
-  URIToFileCache(PathRef HintPath) : HintPath(HintPath) {}
-
-  llvm::StringRef operator[](llvm::StringRef FileURI) {
-    if (FileURI.empty())
-      return "";
-    auto I = URIToPathCache.try_emplace(FileURI);
-    if (I.second) {
-      I.first->second = llvm::cantFail(URI::resolve(FileURI, HintPath),
-                                       "Failed to resolve URI");
-    }
-    return I.first->second;
-  }
-
-private:
-  PathRef HintPath;
-  llvm::StringMap<std::string> URIToPathCache;
-};
-
 // We keep only the node "U" and its edges. Any node other than "U" will be
 // empty in the resultant graph.
 IncludeGraph getSubGraph(llvm::StringRef URI, const IncludeGraph &FullGraph) {
   IncludeGraph IG;
 
   auto Entry = IG.try_emplace(URI).first;
   auto &Node = Entry->getValue();
   Node = FullGraph.lookup(Entry->getKey());
   Node.URI = Entry->getKey();
 
   // URIs inside nodes must point into the keys of the same IncludeGraph.
   for (auto &Include : Node.DirectIncludes) {
     auto I = IG.try_emplace(Include).first;
     I->getValue().URI = I->getKey();
     Include = I->getKey();
   }
   return IG;
 }
 } // namespace
 
-FileShardedIndex::FileShardedIndex(IndexFileIn Input, PathRef HintPath)
+FileShardedIndex::FileShardedIndex(IndexFileIn Input)
     : Index(std::move(Input)) {
-  URIToFileCache UriToFile(HintPath);
   // Used to build RelationSlabs.
   llvm::DenseMap<SymbolID, FileShard *> SymbolIDToFile;
 
   // Attribute each Symbol to both their declaration and definition locations.
   if (Index.Symbols) {
     for (const auto &S : *Index.Symbols) {
-      auto File = UriToFile[S.CanonicalDeclaration.FileURI];
-      auto It = Shards.try_emplace(File);
+      auto It = Shards.try_emplace(S.CanonicalDeclaration.FileURI);
       It.first->getValue().Symbols.insert(&S);
       SymbolIDToFile[S.ID] = &It.first->getValue();
       // Only bother if definition file is different than declaration file.
       if (S.Definition &&
           S.Definition.FileURI != S.CanonicalDeclaration.FileURI) {
-        auto File = UriToFile[S.Definition.FileURI];
-        auto It = Shards.try_emplace(File);
+        auto It = Shards.try_emplace(S.Definition.FileURI);
         It.first->getValue().Symbols.insert(&S);
       }
     }
   }
   // Attribute references into each file they occured in.
   if (Index.Refs) {
     for (const auto &SymRefs : *Index.Refs) {
       for (const auto &R : SymRefs.second) {
-        auto File = UriToFile[R.Location.FileURI];
-        const auto It = Shards.try_emplace(File);
+        const auto It = Shards.try_emplace(R.Location.FileURI);
         It.first->getValue().Refs.insert(&R);
         RefToSymID[&R] = SymRefs.first;
       }
     }
   }
   // Attribute relations to the file declaraing their Subject as Object might
   // not have been indexed, see SymbolCollector::processRelations for details.
   if (Index.Relations) {
     for (const auto &R : *Index.Relations) {
       // FIXME: RelationSlab shouldn't contain dangling relations.
       if (auto *File = SymbolIDToFile.lookup(R.Subject))
         File->Relations.insert(&R);
     }
   }
   // Store only the direct includes of a file in a shard.
   if (Index.Sources) {
     const auto &FullGraph = *Index.Sources;
     for (const auto &It : FullGraph) {
-      auto File = UriToFile[It.first()];
-      auto ShardIt = Shards.try_emplace(File);
+      auto ShardIt = Shards.try_emplace(It.first());
       ShardIt.first->getValue().IG = getSubGraph(It.first(), FullGraph);
     }
   }
 }
-std::vector<PathRef> FileShardedIndex::getAllFiles() const {
+std::vector<llvm::StringRef> FileShardedIndex::getAllSources() const {
   // It should be enough to construct a vector with {Shards.keys().begin(),
   // Shards.keys().end()} but MSVC fails to compile that.
   std::vector<PathRef> Result;
   Result.reserve(Shards.size());
-  for (PathRef Key : Shards.keys())
+  for (auto Key : Shards.keys())
     Result.push_back(Key);
   return Result;
 }
 
-IndexFileIn FileShardedIndex::getShard(PathRef File) const {
-  auto It = Shards.find(File);
-  assert(It != Shards.end() && "received unknown file");
+llvm::Optional<IndexFileIn>
+FileShardedIndex::getShard(llvm::StringRef Uri) const {
+  auto It = Shards.find(Uri);
+  if (It == Shards.end())
+    return llvm::None;
 
   IndexFileIn IF;
   IF.Sources = It->getValue().IG;
   IF.Cmd = Index.Cmd;
 
   SymbolSlab::Builder SymB;
   for (const auto *S : It->getValue().Symbols)
     SymB.insert(*S);
SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
   std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
   return indexSymbols(AST, std::move(PP), DeclsToIndex,
                       /*MainFileMacros=*/nullptr, Includes,
                       /*IsIndexMainAST=*/false, Version);
 }
 
-void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Symbols,
+void FileSymbols::update(llvm::StringRef Key,
+                         std::unique_ptr<SymbolSlab> Symbols,
                          std::unique_ptr<RefSlab> Refs,
                          std::unique_ptr<RelationSlab> Relations,
                          bool CountReferences) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Symbols)
-    FileToSymbols.erase(Path);
+    SymbolsSnapshot.erase(Key);
   else
-    FileToSymbols[Path] = std::move(Symbols);
+    SymbolsSnapshot[Key] = std::move(Symbols);
   if (!Refs) {
-    FileToRefs.erase(Path);
+    RefsSnapshot.erase(Key);
   } else {
     RefSlabAndCountReferences Item;
     Item.CountReferences = CountReferences;
     Item.Slab = std::move(Refs);
-    FileToRefs[Path] = std::move(Item);
+    RefsSnapshot[Key] = std::move(Item);
   }
   if (!Relations)
-    FileToRelations.erase(Path);
+    RelatiosSnapshot.erase(Key);
   else
-    FileToRelations[Path] = std::move(Relations);
+    RelatiosSnapshot[Key] = std::move(Relations);
 }
 
 std::unique_ptr<SymbolIndex>
 FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle) {
   std::vector<std::shared_ptr<SymbolSlab>> SymbolSlabs;
   std::vector<std::shared_ptr<RefSlab>> RefSlabs;
   std::vector<std::shared_ptr<RelationSlab>> RelationSlabs;
   std::vector<RefSlab *> MainFileRefs;
   {
     std::lock_guard<std::mutex> Lock(Mutex);
-    for (const auto &FileAndSymbols : FileToSymbols)
+    for (const auto &FileAndSymbols : SymbolsSnapshot)
       SymbolSlabs.push_back(FileAndSymbols.second);
-    for (const auto &FileAndRefs : FileToRefs) {
+    for (const auto &FileAndRefs : RefsSnapshot) {
       RefSlabs.push_back(FileAndRefs.second.Slab);
       if (FileAndRefs.second.CountReferences)
         MainFileRefs.push_back(RefSlabs.back().get());
     }
-    for (const auto &FileAndRelations : FileToRelations)
+    for (const auto &FileAndRelations : RelatiosSnapshot)
       RelationSlabs.push_back(FileAndRelations.second);
   }
   std::vector<const Symbol *> AllSymbols;
   std::vector<Symbol> SymsStorage;
   switch (DuplicateHandle) {
   case DuplicateHandling::Merge: {
     llvm::DenseMap<SymbolID, Symbol> Merged;
     for (const auto &Slab : SymbolSlabs) {

 
 void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
                                ASTContext &AST,
                                std::shared_ptr<Preprocessor> PP,
                                const CanonicalIncludes &Includes) {
   IndexFileIn IF;
   std::tie(IF.Symbols, std::ignore, IF.Relations) =
       indexHeaderSymbols(Version, AST, std::move(PP), Includes);
-  FileShardedIndex ShardedIndex(std::move(IF), Path);
-  for (PathRef File : ShardedIndex.getAllFiles()) {
-    auto IF = ShardedIndex.getShard(File);
+  FileShardedIndex ShardedIndex(std::move(IF));
+  for (auto Uri : ShardedIndex.getAllSources()) {
+    // We are using the key received from ShardedIndex, so it should always
+    // exist.
+    auto IF = ShardedIndex.getShard(Uri).getValue();
     PreambleSymbols.update(
-        File, std::make_unique<SymbolSlab>(std::move(*IF.Symbols)),
+        Uri, std::make_unique<SymbolSlab>(std::move(*IF.Symbols)),
         std::make_unique<RefSlab>(),
         std::make_unique<RelationSlab>(std::move(*IF.Relations)),
         /*CountReferences=*/false);
   }
   PreambleIndex.reset(
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
                                  DuplicateHandling::PickOne));
   vlog("Build dynamic index for header symbols with estimated memory usage of "

clang-tools-extra/clangd/unittests/FileIndexTests.cpp

IncludeGraph &IG = *IF.Sources;
     // a.h includes nothing.
     auto &Node = IG[AHeaderUri];
     Node.DirectIncludes = {};
     Node.URI = AHeaderUri;
   }
 
   IF.Cmd = tooling::CompileCommand(testRoot(), "b.cc", {"clang"}, "out");
 
-  FileShardedIndex ShardedIndex(std::move(IF), testPath("b.cc"));
-  ASSERT_THAT(
-      ShardedIndex.getAllFiles(),
-      UnorderedElementsAre(testPath("a.h"), testPath("b.h"), testPath("b.cc")));
+  FileShardedIndex ShardedIndex(std::move(IF));
+  ASSERT_THAT(ShardedIndex.getAllSources(),
+              UnorderedElementsAre(AHeaderUri, BHeaderUri, BSourceUri));
 
   {
-    auto Shard = ShardedIndex.getShard(testPath("a.h"));
+    auto Shard = ShardedIndex.getShard(AHeaderUri).getValue();
     EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("1")));
     EXPECT_THAT(Shard.Refs.getValue(), IsEmpty());
     EXPECT_THAT(
         Shard.Relations.getValue(),
         UnorderedElementsAre(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID}));
     ASSERT_THAT(Shard.Sources.getValue().keys(),
                 UnorderedElementsAre(AHeaderUri));
     EXPECT_THAT(Shard.Sources->lookup(AHeaderUri).DirectIncludes, IsEmpty());
     EXPECT_TRUE(Shard.Cmd.hasValue());
   }
   {
-    auto Shard = ShardedIndex.getShard(testPath("b.h"));
+    auto Shard = ShardedIndex.getShard(BHeaderUri).getValue();
     EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("2")));
     EXPECT_THAT(Shard.Refs.getValue(), IsEmpty());
     EXPECT_THAT(
         Shard.Relations.getValue(),
         UnorderedElementsAre(Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID}));
     ASSERT_THAT(Shard.Sources.getValue().keys(),
                 UnorderedElementsAre(BHeaderUri, AHeaderUri));
     EXPECT_THAT(Shard.Sources->lookup(BHeaderUri).DirectIncludes,
                 UnorderedElementsAre(AHeaderUri));
     EXPECT_TRUE(Shard.Cmd.hasValue());
   }
   {
-    auto Shard = ShardedIndex.getShard(testPath("b.cc"));
+    auto Shard = ShardedIndex.getShard(BSourceUri).getValue();
     EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("2")));
     EXPECT_THAT(Shard.Refs.getValue(), UnorderedElementsAre(Pair(Sym1.ID, _)));
     EXPECT_THAT(Shard.Relations.getValue(), IsEmpty());
     ASSERT_THAT(Shard.Sources.getValue().keys(),
                 UnorderedElementsAre(BSourceUri, BHeaderUri));
     EXPECT_THAT(Shard.Sources->lookup(BSourceUri).DirectIncludes,
                 UnorderedElementsAre(BHeaderUri));
     EXPECT_TRUE(Shard.Cmd.hasValue());
   }
 }
 } // namespace
 } // namespace clangd
 } // namespace clang