[clangd] Boost completion score according to file proximity.

Summary:
Also move unittest: URI scheme to TestFS so that it can be shared by
different tests.

Reviewers: sammccall

Reviewed By: sammccall

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

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

4 files changed, 175 insertions, 12 deletions

diff --git a/clang-tools-extra/clangd/CodeComplete.cpp b/clang-tools-extra/clangd/CodeComplete.cpp
index a5809c67922..0bec281a388 100644
--- a/clang-tools-extra/clangd/CodeComplete.cpp
+++ b/clang-tools-extra/clangd/CodeComplete.cpp
@@ -848,11 +848,17 @@ class CodeCompleteFlow {
   bool Incomplete = false; // Would more be available with a higher limit?
   llvm::Optional<FuzzyMatcher> Filter;       // Initialized once Sema runs.
   std::unique_ptr<IncludeInserter> Includes; // Initialized once compiler runs.
+  FileProximityMatcher FileProximityMatch;
 
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
   CodeCompleteFlow(PathRef FileName, const CodeCompleteOptions &Opts)
-      : FileName(FileName), Opts(Opts) {}
+      : FileName(FileName), Opts(Opts),
+        // FIXME: also use path of the main header corresponding to FileName to
+        // calculate the file proximity, which would capture include/ and src/
+        // project setup where headers and implementations are not in the same
+        // directory.
+        FileProximityMatch({FileName}) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     trace::Span Tracer("CodeCompleteFlow");
@@ -993,6 +999,7 @@ private:
     SymbolQualitySignals Quality;
     SymbolRelevanceSignals Relevance;
     Relevance.Query = SymbolRelevanceSignals::CodeComplete;
+    Relevance.FileProximityMatch = &FileProximityMatch;
     if (auto FuzzyScore = fuzzyScore(C))
       Relevance.NameMatch = *FuzzyScore;
     else
diff --git a/clang-tools-extra/clangd/Quality.cpp b/clang-tools-extra/clangd/Quality.cpp
index fb67c3a9643..1799a5329c6 100644
--- a/clang-tools-extra/clangd/Quality.cpp
+++ b/clang-tools-extra/clangd/Quality.cpp
@@ -7,6 +7,7 @@
 //
 //===---------------------------------------------------------------------===//
 #include "Quality.h"
+#include "URI.h"
 #include "index/Index.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/CharInfo.h"
@@ -185,6 +186,60 @@ raw_ostream &operator<<(raw_ostream &OS, const SymbolQualitySignals &S) {
   return OS;
 }
 
+/// Calculates a proximity score from \p From and \p To, which are URI strings
+/// that have the same scheme. This does not parse URI. A URI (sans "<scheme>:")
+/// is split into chunks by '/' and each chunk is considered a file/directory.
+/// For example, "uri:///a/b/c" will be treated as /a/b/c
+static float uriProximity(StringRef From, StringRef To) {
+  auto SchemeSplitFrom = From.split(':');
+  auto SchemeSplitTo = To.split(':');
+  assert((SchemeSplitFrom.first == SchemeSplitTo.first) &&
+         "URIs must have the same scheme in order to compute proximity.");
+  auto Split = [](StringRef URIWithoutScheme) {
+    SmallVector<StringRef, 8> Split;
+    URIWithoutScheme.split(Split, '/', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+    return Split;
+  };
+  SmallVector<StringRef, 8> Fs = Split(SchemeSplitFrom.second);
+  SmallVector<StringRef, 8> Ts = Split(SchemeSplitTo.second);
+  auto F = Fs.begin(), T = Ts.begin(), FE = Fs.end(), TE = Ts.end();
+  for (; F != FE && T != TE && *F == *T; ++F, ++T) {
+  }
+  // We penalize for traversing up and down from \p From to \p To but penalize
+  // less for traversing down because subprojects are more closely related than
+  // superprojects.
+  int UpDist = FE - F;
+  int DownDist = TE - T;
+  return std::pow(0.7, UpDist + DownDist/2);
+}
+
+FileProximityMatcher::FileProximityMatcher(ArrayRef<StringRef> ProximityPaths)
+    : ProximityPaths(ProximityPaths.begin(), ProximityPaths.end()) {}
+
+float FileProximityMatcher::uriProximity(StringRef SymbolURI) const {
+  float Score = 0;
+  if (!ProximityPaths.empty() && !SymbolURI.empty()) {
+    for (const auto &Path : ProximityPaths)
+      // Only calculate proximity score for two URIs with the same scheme so
+      // that the computation can be purely text-based and thus avoid expensive
+      // URI encoding/decoding.
+      if (auto U = URI::create(Path, SymbolURI.split(':').first)) {
+        Score = std::max(Score, clangd::uriProximity(U->toString(), SymbolURI));
+      } else {
+        llvm::consumeError(U.takeError());
+      }
+  }
+  return Score;
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
+                              const FileProximityMatcher &M) {
+  OS << formatv("File proximity matcher: ");
+  OS << formatv("ProximityPaths{{0}}", llvm::join(M.ProximityPaths.begin(),
+                                                  M.ProximityPaths.end(), ","));
+  return OS;
+}
+
 static SymbolRelevanceSignals::AccessibleScope
 ComputeScope(const NamedDecl &D) {
   bool InClass = false;
@@ -205,6 +260,8 @@ ComputeScope(const NamedDecl &D) {
 void SymbolRelevanceSignals::merge(const Symbol &IndexResult) {
   // FIXME: Index results always assumed to be at global scope. If Scope becomes
   // relevant to non-completion requests, we should recognize class members etc.
+
+  SymbolURI = IndexResult.CanonicalDeclaration.FileURI;
 }
 
 void SymbolRelevanceSignals::merge(const CodeCompletionResult &SemaCCResult) {
@@ -213,11 +270,12 @@ void SymbolRelevanceSignals::merge(const CodeCompletionResult &SemaCCResult) {
     Forbidden = true;
 
   if (SemaCCResult.Declaration) {
-    // We boost things that have decls in the main file.
-    // The real proximity scores would be more general when we have them.
+    // We boost things that have decls in the main file. We give a fixed score
+    // for all other declarations in sema as they are already included in the
+    // translation unit.
     float DeclProximity =
-        hasDeclInMainFile(*SemaCCResult.Declaration) ? 1.0 : 0.0;
-    ProximityScore = std::max(DeclProximity, ProximityScore);
+        hasDeclInMainFile(*SemaCCResult.Declaration) ? 1.0 : 0.6;
+    SemaProximityScore = std::max(DeclProximity, SemaProximityScore);
   }
 
   // Declarations are scoped, others (like macros) are assumed global.
@@ -233,9 +291,11 @@ float SymbolRelevanceSignals::evaluate() const {
 
   Score *= NameMatch;
 
+  float IndexProximityScore =
+      FileProximityMatch ? FileProximityMatch->uriProximity(SymbolURI) : 0;
   // Proximity scores are [0,1] and we translate them into a multiplier in the
   // range from 1 to 2.
-  Score *= 1 + ProximityScore;
+  Score *= 1 + std::max(IndexProximityScore, SemaProximityScore);
 
   // Symbols like local variables may only be referenced within their scope.
   // Conversely if we're in that scope, it's likely we'll reference them.
@@ -259,11 +319,18 @@ float SymbolRelevanceSignals::evaluate() const {
 
   return Score;
 }
+
 raw_ostream &operator<<(raw_ostream &OS, const SymbolRelevanceSignals &S) {
   OS << formatv("=== Symbol relevance: {0}\n", S.evaluate());
   OS << formatv("\tName match: {0}\n", S.NameMatch);
   OS << formatv("\tForbidden: {0}\n", S.Forbidden);
-  OS << formatv("\tProximity: {0}\n", S.ProximityScore);
+  OS << formatv("\tSymbol URI: {0}\n", S.SymbolURI);
+  if (S.FileProximityMatch) {
+    OS << formatv("\tIndex proximity: {0}\n",
+                  S.FileProximityMatch->uriProximity(S.SymbolURI))
+       << " (" << *S.FileProximityMatch << ")\n";
+  }
+  OS << formatv("\tSema proximity: {0}\n", S.SemaProximityScore);
   OS << formatv("\tQuery type: {0}\n", static_cast<int>(S.Query));
   OS << formatv("\tScope: {0}\n", static_cast<int>(S.Scope));
   return OS;
diff --git a/clang-tools-extra/clangd/Quality.h b/clang-tools-extra/clangd/Quality.h
index 0e0b6df6ba8..fef039a3a2f 100644
--- a/clang-tools-extra/clangd/Quality.h
+++ b/clang-tools-extra/clangd/Quality.h
@@ -26,6 +26,7 @@
 //===---------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_QUALITY_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_QUALITY_H
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include <algorithm>
 #include <functional>
@@ -68,13 +69,22 @@ struct SymbolQualitySignals {
 llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                               const SymbolQualitySignals &);
 
+class FileProximityMatcher;
+
 /// Attributes of a symbol-query pair that affect how much we like it.
 struct SymbolRelevanceSignals {
   /// 0-1+ fuzzy-match score for unqualified name. Must be explicitly assigned.
   float NameMatch = 1;
   bool Forbidden = false; // Unavailable (e.g const) or inaccessible (private).
+
+  const FileProximityMatcher *FileProximityMatch = nullptr;
+  /// This is used to calculate proximity between the index symbol and the
+  /// query.
+  llvm::StringRef SymbolURI;
   /// Proximity between best declaration and the query. [0-1], 1 is closest.
-  float ProximityScore = 0;
+  /// FIXME: unify with index proximity score - signals should be
+  /// source-independent.
+  float SemaProximityScore = 0;
 
   // An approximate measure of where we expect the symbol to be used.
   enum AccessibleScope {
@@ -101,6 +111,25 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &,
 /// Combine symbol quality and relevance into a single score.
 float evaluateSymbolAndRelevance(float SymbolQuality, float SymbolRelevance);
 
+class FileProximityMatcher {
+  public:
+    /// \p ProximityPaths are used to compute proximity scores from symbol's
+    /// declaring file. The best score will be used.
+    explicit FileProximityMatcher(
+        llvm::ArrayRef<llvm::StringRef> ProximityPaths);
+
+    /// Calculates the best proximity score from proximity paths to the symbol's
+    /// URI. Score is [0-1], 1 means \p SymbolURI exactly matches a proximity
+    /// path. When a path cannot be encoded into the same scheme as \p
+    /// SymbolURI, the proximity will be 0.
+    float uriProximity(llvm::StringRef SymbolURI) const;
+
+  private:
+    llvm::SmallVector<std::string, 2> ProximityPaths;
+    friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
+                                         const FileProximityMatcher &);
+};
+
 /// TopN<T> is a lossy container that preserves only the "best" N elements.
 template <typename T, typename Compare = std::greater<T>> class TopN {
 public:
diff --git a/clang-tools-extra/unittests/clangd/QualityTests.cpp b/clang-tools-extra/unittests/clangd/QualityTests.cpp
index 7238bf68fa7..010a30c24d9 100644
--- a/clang-tools-extra/unittests/clangd/QualityTests.cpp
+++ b/clang-tools-extra/unittests/clangd/QualityTests.cpp
@@ -18,12 +18,18 @@
 //===----------------------------------------------------------------------===//
 
 #include "Quality.h"
+#include "TestFS.h"
 #include "TestTU.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
+
+// Force the unittest URI scheme to be linked,
+static int LLVM_ATTRIBUTE_UNUSED UnittestSchemeAnchorDest =
+    UnittestSchemeAnchorSource;
+
 namespace {
 
 TEST(QualityTests, SymbolQualitySignalExtraction) {
@@ -91,13 +97,14 @@ TEST(QualityTests, SymbolRelevanceSignalExtraction) {
 
   Relevance = {};
   Relevance.merge(CodeCompletionResult(&findDecl(AST, "main"), 42));
-  EXPECT_FLOAT_EQ(Relevance.ProximityScore, 1.0) << "Decl in current file";
+  EXPECT_FLOAT_EQ(Relevance.SemaProximityScore, 1.0) << "Decl in current file";
   Relevance = {};
   Relevance.merge(CodeCompletionResult(&findDecl(AST, "header"), 42));
-  EXPECT_FLOAT_EQ(Relevance.ProximityScore, 0.0) << "Decl from header";
+  EXPECT_FLOAT_EQ(Relevance.SemaProximityScore, 0.6) << "Decl from header";
   Relevance = {};
   Relevance.merge(CodeCompletionResult(&findDecl(AST, "header_main"), 42));
-  EXPECT_FLOAT_EQ(Relevance.ProximityScore, 1.0) << "Current file and header";
+  EXPECT_FLOAT_EQ(Relevance.SemaProximityScore, 1.0)
+      << "Current file and header";
 
   Relevance = {};
   Relevance.merge(CodeCompletionResult(&findAnyDecl(AST, "X"), 42));
@@ -151,7 +158,7 @@ TEST(QualityTests, SymbolRelevanceSignalsSanity) {
   EXPECT_LT(PoorNameMatch.evaluate(), Default.evaluate());
 
   SymbolRelevanceSignals WithProximity;
-  WithProximity.ProximityScore = 0.2f;
+  WithProximity.SemaProximityScore = 0.2f;
   EXPECT_GT(WithProximity.evaluate(), Default.evaluate());
 
   SymbolRelevanceSignals Scoped;
@@ -173,6 +180,59 @@ TEST(QualityTests, SortText) {
   EXPECT_LT(sortText(0, "a"), sortText(0, "z"));
 }
 
+// {a,b,c} becomes /clangd-test/a/b/c
+std::string joinPaths(llvm::ArrayRef<StringRef> Parts) {
+  return testPath(
+      llvm::join(Parts.begin(), Parts.end(), llvm::sys::path::get_separator()));
+}
+
+static constexpr float ProximityBase = 0.7;
+
+// Calculates a proximity score for an index symbol with declaration file
+// SymPath with the given URI scheme.
+float URIProximity(const FileProximityMatcher &Matcher, StringRef SymPath,
+                     StringRef Scheme = "file") {
+  auto U = URI::create(SymPath, Scheme);
+  EXPECT_TRUE(static_cast<bool>(U)) << llvm::toString(U.takeError());
+  return Matcher.uriProximity(U->toString());
+}
+
+TEST(QualityTests, URIProximityScores) {
+  FileProximityMatcher Matcher(
+      /*ProximityPaths=*/{joinPaths({"a", "b", "c", "d", "x"})});
+
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"a", "b", "c", "d", "x"})),
+                  1);
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"a", "b", "c", "d", "y"})),
+                  ProximityBase);
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"a", "y", "z"})),
+                  std::pow(ProximityBase, 5));
+  EXPECT_FLOAT_EQ(
+      URIProximity(Matcher, joinPaths({"a", "b", "c", "d", "e", "y"})),
+      std::pow(ProximityBase, 2));
+  EXPECT_FLOAT_EQ(
+      URIProximity(Matcher, joinPaths({"a", "b", "m", "n", "o", "y"})),
+      std::pow(ProximityBase, 5));
+  EXPECT_FLOAT_EQ(
+      URIProximity(Matcher, joinPaths({"a", "t", "m", "n", "o", "y"})),
+      std::pow(ProximityBase, 6));
+  // Note the common directory is /clang-test/
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"m", "n", "o", "p", "y"})),
+                  std::pow(ProximityBase, 7));
+}
+
+TEST(QualityTests, URIProximityScoresWithTestURI) {
+  FileProximityMatcher Matcher(
+      /*ProximityPaths=*/{joinPaths({"b", "c", "x"})});
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"b", "c", "x"}), "unittest"),
+                  1);
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"b", "y"}), "unittest"),
+                  std::pow(ProximityBase, 2));
+  // unittest:///b/c/x vs unittest:///m/n/y. No common directory.
+  EXPECT_FLOAT_EQ(URIProximity(Matcher, joinPaths({"m", "n", "y"}), "unittest"),
+                  std::pow(ProximityBase, 4));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang