[clangd] Add type boost to fuzzy find in Dex.

Summary:
No noticeable impact on code completions overall except some improvement on
cross-namespace completion.

Reviewers: sammccall, ilya-biryukov

Reviewed By: sammccall

Subscribers: MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

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

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@353310 91177308-0d34-0410-b5e6-96231b3b80d8

7 files changed, 97 insertions, 50 deletions

diff --git a/clangd/CodeComplete.cpp b/clangd/CodeComplete.cpp
index 83dcfb7a..b059216a 100644
--- a/clangd/CodeComplete.cpp
+++ b/clangd/CodeComplete.cpp
@@ -1333,6 +1333,8 @@ private:
     Req.AnyScope = AllScopes;
     // FIXME: we should send multiple weighted paths here.
     Req.ProximityPaths.push_back(FileName);
+    if (PreferredType)
+      Req.PreferredTypes.push_back(PreferredType->raw());
     vlog("Code complete: fuzzyFind({0:2})", toJSON(Req));
 
     if (SpecFuzzyFind)
diff --git a/clangd/index/Index.cpp b/clangd/index/Index.cpp
index c773cd69..36c591d9 100644
--- a/clangd/index/Index.cpp
+++ b/clangd/index/Index.cpp
@@ -179,7 +179,8 @@ bool fromJSON(const llvm::json::Value &Parameters, FuzzyFindRequest &Request) {
       O && O.map("Query", Request.Query) && O.map("Scopes", Request.Scopes) &&
       O.map("AnyScope", Request.AnyScope) && O.map("Limit", Limit) &&
       O.map("RestrictForCodeCompletion", Request.RestrictForCodeCompletion) &&
-      O.map("ProximityPaths", Request.ProximityPaths);
+      O.map("ProximityPaths", Request.ProximityPaths) &&
+      O.map("PreferredTypes", Request.PreferredTypes);
   if (OK && Limit <= std::numeric_limits<uint32_t>::max())
     Request.Limit = Limit;
   return OK;
@@ -193,6 +194,7 @@ llvm::json::Value toJSON(const FuzzyFindRequest &Request) {
       {"Limit", Request.Limit},
       {"RestrictForCodeCompletion", Request.RestrictForCodeCompletion},
       {"ProximityPaths", llvm::json::Array{Request.ProximityPaths}},
+      {"PreferredTypes", llvm::json::Array{Request.PreferredTypes}},
   };
 }
 
diff --git a/clangd/index/Index.h b/clangd/index/Index.h
index d36dd299..4f6feb93 100644
--- a/clangd/index/Index.h
+++ b/clangd/index/Index.h
@@ -454,14 +454,15 @@ struct FuzzyFindRequest {
   /// Contextually relevant files (e.g. the file we're code-completing in).
   /// Paths should be absolute.
   std::vector<std::string> ProximityPaths;
-
-  // FIXME(ibiryukov): add expected type to the request.
+  /// Preferred types of symbols. These are raw representation of `OpaqueType`.
+  std::vector<std::string> PreferredTypes;
 
   bool operator==(const FuzzyFindRequest &Req) const {
     return std::tie(Query, Scopes, Limit, RestrictForCodeCompletion,
-                    ProximityPaths) ==
+                    ProximityPaths, PreferredTypes) ==
            std::tie(Req.Query, Req.Scopes, Req.Limit,
-                    Req.RestrictForCodeCompletion, Req.ProximityPaths);
+                    Req.RestrictForCodeCompletion, Req.ProximityPaths,
+                    Req.PreferredTypes);
   }
   bool operator!=(const FuzzyFindRequest &Req) const { return !(*this == Req); }
 };
diff --git a/clangd/index/dex/Dex.cpp b/clangd/index/dex/Dex.cpp
index b4db1136..d767bb51 100644
--- a/clangd/index/dex/Dex.cpp
+++ b/clangd/index/dex/Dex.cpp
@@ -42,7 +42,6 @@ const Token RestrictedForCodeCompletion =
 // Returns the tokens which are given symbols's characteristics. For example,
 // trigrams and scopes.
 // FIXME(kbobyrev): Support more token types:
-// * Types
 // * Namespace proximity
 std::vector<Token> generateSearchTokens(const Symbol &Sym) {
   std::vector<Token> Result = generateIdentifierTrigrams(Sym.Name);
@@ -54,49 +53,11 @@ std::vector<Token> generateSearchTokens(const Symbol &Sym) {
       Result.emplace_back(Token::Kind::ProximityURI, ProximityURI);
   if (Sym.Flags & Symbol::IndexedForCodeCompletion)
     Result.emplace_back(RestrictedForCodeCompletion);
+  if (!Sym.Type.empty())
+    Result.emplace_back(Token::Kind::Type, Sym.Type);
   return Result;
 }
 
-// Constructs BOOST iterators for Path Proximities.
-std::unique_ptr<Iterator> createFileProximityIterator(
-    llvm::ArrayRef<std::string> ProximityPaths,
-    const llvm::DenseMap<Token, PostingList> &InvertedIndex,
-    const Corpus &Corpus) {
-  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());
-    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.
-  URIDistance DistanceCalculator(Sources);
-  PathProximitySignals.FileProximityMatch = &DistanceCalculator;
-  // Try to build BOOST iterator for each Proximity Path provided by
-  // ProximityPaths. Boosting factor should depend on the distance to the
-  // Proximity Path: the closer processed path is, the higher boosting factor.
-  for (const auto &ParentURI : ParentURIs.keys()) {
-    Token Tok(Token::Kind::ProximityURI, ParentURI);
-    const auto It = InvertedIndex.find(Tok);
-    if (It != InvertedIndex.end()) {
-      // FIXME(kbobyrev): Append LIMIT on top of every BOOST iterator.
-      PathProximitySignals.SymbolURI = ParentURI;
-      BoostingIterators.push_back(Corpus.boost(
-          It->second.iterator(&It->first), PathProximitySignals.evaluate()));
-    }
-  }
-  BoostingIterators.push_back(Corpus.all());
-  return Corpus.unionOf(std::move(BoostingIterators));
-}
-
 } // namespace
 
 void Dex::buildIndex() {
@@ -141,6 +102,57 @@ std::unique_ptr<Iterator> Dex::iterator(const Token &Tok) const {
                                    : It->second.iterator(&It->first);
 }
 
+// Constructs BOOST iterators for Path Proximities.
+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());
+    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.
+  URIDistance DistanceCalculator(Sources);
+  PathProximitySignals.FileProximityMatch = &DistanceCalculator;
+  // Try to build BOOST iterator for each Proximity Path provided by
+  // ProximityPaths. Boosting factor should depend on the distance to the
+  // Proximity Path: the closer processed path is, the higher boosting factor.
+  for (const auto &ParentURI : ParentURIs.keys()) {
+    // FIXME(kbobyrev): Append LIMIT on top of every BOOST iterator.
+    auto It = iterator(Token(Token::Kind::ProximityURI, ParentURI));
+    if (It->kind() != Iterator::Kind::False) {
+      PathProximitySignals.SymbolURI = ParentURI;
+      BoostingIterators.push_back(
+          Corpus.boost(std::move(It), PathProximitySignals.evaluate()));
+    }
+  }
+  BoostingIterators.push_back(Corpus.all());
+  return Corpus.unionOf(std::move(BoostingIterators));
+}
+
+// Constructs BOOST iterators for preferred types.
+std::unique_ptr<Iterator>
+Dex::createTypeBoostingIterator(llvm::ArrayRef<std::string> Types) const {
+  std::vector<std::unique_ptr<Iterator>> BoostingIterators;
+  SymbolRelevanceSignals PreferredTypeSignals;
+  PreferredTypeSignals.TypeMatchesPreferred = true;
+  auto Boost = PreferredTypeSignals.evaluate();
+  for (const auto &T : Types)
+    BoostingIterators.push_back(
+        Corpus.boost(iterator(Token(Token::Kind::Type, T)), Boost));
+  BoostingIterators.push_back(Corpus.all());
+  return Corpus.unionOf(std::move(BoostingIterators));
+}
+
 /// Constructs iterators over tokens extracted from the query and exhausts it
 /// while applying Callback to each symbol in the order of decreasing quality
 /// of the matched symbols.
@@ -174,8 +186,9 @@ bool Dex::fuzzyFind(const FuzzyFindRequest &Req,
   Criteria.push_back(Corpus.unionOf(move(ScopeIterators)));
 
   // Add proximity paths boosting (all symbols, some boosted).
-  Criteria.push_back(
-      createFileProximityIterator(Req.ProximityPaths, InvertedIndex, Corpus));
+  Criteria.push_back(createFileProximityIterator(Req.ProximityPaths));
+  // Add boosting for preferred types.
+  Criteria.push_back(createTypeBoostingIterator(Req.PreferredTypes));
 
   if (Req.RestrictForCodeCompletion)
     Criteria.push_back(iterator(RestrictedForCodeCompletion));
diff --git a/clangd/index/dex/Dex.h b/clangd/index/dex/Dex.h
index 71814e43..fb80ca03 100644
--- a/clangd/index/dex/Dex.h
+++ b/clangd/index/dex/Dex.h
@@ -77,6 +77,10 @@ public:
 private:
   void buildIndex();
   std::unique_ptr<Iterator> iterator(const Token &Tok) const;
+  std::unique_ptr<Iterator>
+  createFileProximityIterator(llvm::ArrayRef<std::string> ProximityPaths) const;
+  std::unique_ptr<Iterator>
+  createTypeBoostingIterator(llvm::ArrayRef<std::string> Types) const;
 
   /// Stores symbols sorted in the descending order of symbol quality..
   std::vector<const Symbol *> Symbols;
diff --git a/clangd/index/dex/Token.h b/clangd/index/dex/Token.h
index 4c67e9fa..37859bcf 100644
--- a/clangd/index/dex/Token.h
+++ b/clangd/index/dex/Token.h
@@ -62,11 +62,11 @@ struct Token {
     /// Example: "file:///path/to/clang-tools-extra/clangd/index/SymbolIndex.h"
     /// and some amount of its parents.
     ProximityURI,
+    /// Type of symbol (see `Symbol::Type`).
+    Type,
     /// Internal Token type for invalid/special tokens, e.g. empty tokens for
     /// llvm::DenseMap.
     Sentinel,
-    /// FIXME(kbobyrev): Add other Token Kinds
-    /// * Type with qualified type name or its USR
   };
 
   Token(Kind TokenKind, llvm::StringRef Data)
@@ -91,6 +91,9 @@ struct Token {
     case Kind::ProximityURI:
       OS << "U=";
       break;
+    case Kind::Type:
+      OS << "Ty=";
+      break;
     case Kind::Sentinel:
       OS << "?=";
       break;
diff --git a/unittests/clangd/DexTests.cpp b/unittests/clangd/DexTests.cpp
index ce318e2e..cd52daac 100644
--- a/unittests/clangd/DexTests.cpp
+++ b/unittests/clangd/DexTests.cpp
@@ -688,6 +688,28 @@ TEST(DexTests, Refs) {
   EXPECT_THAT(Files, ElementsAre(AnyOf("foo.h", "foo.cc")));
 }
 
+TEST(DexTest, PreferredTypesBoosting) {
+  auto Sym1 = symbol("t1");
+  Sym1.Type = "T1";
+  auto Sym2 = symbol("t2");
+  Sym2.Type = "T2";
+
+  std::vector<Symbol> Symbols{Sym1, Sym2};
+  Dex I(Symbols, RefSlab());
+
+  FuzzyFindRequest Req;
+  Req.AnyScope = true;
+  Req.Query = "t";
+  // The best candidate can change depending on the preferred type.
+  Req.Limit = 1;
+
+  Req.PreferredTypes = {Sym1.Type};
+  EXPECT_THAT(match(I, Req), ElementsAre("t1"));
+
+  Req.PreferredTypes = {Sym2.Type};
+  EXPECT_THAT(match(I, Req), ElementsAre("t2"));
+}
+
 } // namespace
 } // namespace dex
 } // namespace clangd