[clangd] Optimize Dex::generateProximityURIs().

[sammccall]

Summary

Production profiles show that generateProximityURIs is roughly 3.8% of
buildPreamble. Of this, the majority (3% of buildPreamble) is parsing
and reserializing URIs.

We can do this with ugly string manipulation instead.

[adamcz]

clang-tools-extra/clangd/index/dex/Dex.cpp

358

Is there a reason why you're doing this with manual manipulation of C-strings instead of StringRefs? I suspect passing StringRef here was not the problem, turning it into an std::string, from another std::string inside ParsedURI was.

Something like:
S = S.drop_while([](char c) { return c == ':'; });
S.consume_front(":");
if (S.consume_front("//")) {

S = S.drop_while([](char c) { return c == '/'; });
S.consume_front("/");

}

return S;

would be a bit more readable and less likely to have an off-by-one somewhere, imho.

[sammccall]

pointers -> StringRef

clang-tools-extra/clangd/index/dex/Dex.cpp

358

I started with StringRef, and found it too awkward. I've changed it back, but I think it makes the main algorithm harder to understand (it's OK for findPathInURI, but mixing styles is more confusing for me).

StringRef is good at modelling range content or two pointers, but this algorithm is fundamentally about three pointers and uses (Begin, End) and (Path, End) equally. So you end up with too many degrees of freedom in the code and have to paper over with assertions.
And StringRef's API is mostly suited for manipulating content, not dealing with overlapping refs, so you end up with a mix of index+pointer+stringref that personally I find harder to understand.

I'm probably an outlier here, so converted it to use StringRef anyway

clang-tools-extra/clangd/index/dex/Dex.h

private:
   // Size of memory retained by KeepAlive.
   size_t BackingDataSize = 0;
 };
 
 /// Returns Search Token for a number of parent directories of given Path.
 /// Should be used within the index build process.
 ///
 /// This function is exposed for testing only.
-std::vector<std::string> generateProximityURIs(llvm::StringRef URIPath);
+llvm::SmallVector<llvm::StringRef, 5> generateProximityURIs(llvm::StringRef);
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_DEX_DEX_H

clang-tools-extra/clangd/index/dex/Dex.cpp

 std::unique_ptr<Iterator> Dex::createFileProximityIterator(
     llvm::ArrayRef<std::string> ProximityPaths) const {
   std::vector<std::unique_ptr<Iterator>> BoostingIterators;
   // Deduplicate parent URIs extracted from the ProximityPaths.
   llvm::StringSet<> ParentURIs;
   llvm::StringMap<SourceParams> Sources;
   for (const auto &Path : ProximityPaths) {
     Sources[Path] = SourceParams();
-    auto PathURI = URI::create(Path);
-    const auto PathProximityURIs = generateProximityURIs(PathURI.toString());
+    auto PathURI = URI::create(Path).toString();
+    const auto PathProximityURIs = generateProximityURIs(PathURI.c_str());
     for (const auto &ProximityURI : PathProximityURIs)
       ParentURIs.insert(ProximityURI);
   }
   // Use SymbolRelevanceSignals for symbol relevance evaluation: use defaults
   // for all parameters except for Proximity Path distance signal.
   SymbolRelevanceSignals PathProximitySignals;
   // DistanceCalculator will find the shortest distance from ProximityPaths to
   // any URI extracted from the ProximityPaths.
size_t Dex::estimateMemoryUsage() const {
   Bytes += InvertedIndex.getMemorySize();
   for (const auto &TokenToPostingList : InvertedIndex)
     Bytes += TokenToPostingList.second.bytes();
   Bytes += Refs.getMemorySize();
   Bytes += Relations.getMemorySize();
   return Bytes + BackingDataSize;
 }
 
-std::vector<std::string> generateProximityURIs(llvm::StringRef URIPath) {
-  std::vector<std::string> Result;
-  auto ParsedURI = URI::parse(URIPath);
-  assert(ParsedURI &&
-         "Non-empty argument of generateProximityURIs() should be a valid "
-         "URI.");
-  llvm::StringRef Body = ParsedURI->body();
-  // FIXME(kbobyrev): Currently, this is a heuristic which defines the maximum
-  // size of resulting vector. Some projects might want to have higher limit if
-  // the file hierarchy is deeper. For the generic case, it would be useful to
-  // calculate Limit in the index build stage by calculating the maximum depth
-  // of the project source tree at runtime.
-  size_t Limit = 5;
-  // Insert original URI before the loop: this would save a redundant iteration
-  // with a URI parse.
-  Result.emplace_back(ParsedURI->toString());
-  while (!Body.empty() && --Limit > 0) {
-    // FIXME(kbobyrev): Parsing and encoding path to URIs is not necessary and
-    // could be optimized.
-    Body = llvm::sys::path::parent_path(Body, llvm::sys::path::Style::posix);
-    if (!Body.empty())
-      Result.emplace_back(
-          URI(ParsedURI->scheme(), ParsedURI->authority(), Body).toString());
+// Given foo://bar/one/two
+// Returns        ~~~~~~~~  (or empty for bad URI)
+llvm::StringRef findPathInURI(llvm::StringRef S) {
+  S = S.split(':').second;   // Eat scheme.
+  if (S.consume_front("//")) // Eat authority.
+    S = S.drop_until([](char C) { return C == '/'; });
+  return S;
+}
+
+constexpr unsigned ProximityURILimit = 5;
+
+llvm::SmallVector<llvm::StringRef, ProximityURILimit>
+generateProximityURIs(llvm::StringRef URI) {
+  // This function is hot when indexing, so don't parse/reserialize URIPath,
+  // just emit substrings of it instead.
+  //
+  // foo://bar/one/two
+  //          ~~~~~~~~
+  //            Path
+  llvm::StringRef Path = findPathInURI(URI);
+  if (Path.empty())
+    return {}; // Bad URI.
+  assert(Path.begin() >= URI.begin() && Path.begin() < URI.end() &&
+         Path.end() == URI.end());
+
+  // The original is a proximity URI.
+  llvm::SmallVector<llvm::StringRef, ProximityURILimit> Result = {URI};
+  // Form proximity URIs by chopping before each slash in the path part.
+  for (auto Slash = Path.rfind('/'); Slash > 0 && Slash != StringRef::npos;
+       Slash = Path.rfind('/')) {
+    Path = Path.substr(0, Slash);
+    Result.push_back(URI.substr(0, Path.end() - URI.data()));
+    if (Result.size() == ProximityURILimit)
+      return Result;
   }
+  // The root foo://bar/ is a proximity URI.
+  if (Path.startswith("/"))
+    Result.push_back(URI.substr(0, Path.begin() + 1 - URI.data()));
   return Result;
 }
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang