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

[ioeric]

Summary

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

[sammccall]

clangd/index/Index.h

457–458

should these be OpaqueType?

I think the only reason we use StringRef in Symbol is for the non-owning aspect.

otherwise indicate encoding in comment.

clangd/index/dex/Dex.cpp

72

BoostingIterators.push_back(Corpus.boost(iterator(Tok)))?

unittests/clangd/DexTests.cpp

703

proximity paths?

[ioeric]

• address comments plus some refactoring.

[ioeric]

• revert unintended change.

clangd/CodeComplete.cpp

SymbolSlab queryIndex() {
     if (Opts.Limit)
       Req.Limit = Opts.Limit;
     Req.Query = Filter->pattern();
     Req.RestrictForCodeCompletion = true;
     Req.Scopes = QueryScopes;
     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)
       SpecFuzzyFind->NewReq = Req;
     if (SpecFuzzyFind && SpecFuzzyFind->Result.valid() && (*SpecReq == Req)) {
       vlog("Code complete: speculative fuzzy request matches the actual index "
            "request. Waiting for the speculative index results.");
       SPAN_ATTACH(Tracer, "Speculative results", true);

clangd/ExpectedTypes.h

 class OpaqueType {
 public:
   /// Create a type from a code completion result.
   static llvm::Optional<OpaqueType>
   fromCompletionResult(ASTContext &Ctx, const CodeCompletionResult &R);
   /// Construct an instance from a clang::QualType. This is usually a
   /// PreferredType from a clang's completion context.
   static llvm::Optional<OpaqueType> fromType(ASTContext &Ctx, QualType Type);
+  static OpaqueType fromRaw(llvm::StringRef Data) { return OpaqueType(Data); }
 
   /// Get the raw byte representation of the type. You can only rely on the
   /// types being equal iff their raw representation is the same. The particular
   /// details of the used encoding might change over time and one should not
   /// rely on it.
   llvm::StringRef raw() const { return Data; }
 
   friend bool operator==(const OpaqueType &L, const OpaqueType &R) {

clangd/index/Index.h

struct FuzzyFindRequest {
   /// \brief The number of top candidates to return. The index may choose to
   /// return more than this, e.g. if it doesn't know which candidates are best.
   llvm::Optional<uint32_t> Limit;
   /// If set to true, only symbols for completion support will be considered.
   bool RestrictForCodeCompletion = false;
   /// 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); }
 };
 bool fromJSON(const llvm::json::Value &Value, FuzzyFindRequest &Request);
 llvm::json::Value toJSON(const FuzzyFindRequest &Request);
 
 struct LookupRequest {
   llvm::DenseSet<SymbolID> IDs;

clangd/index/Index.cpp

 
 bool fromJSON(const llvm::json::Value &Parameters, FuzzyFindRequest &Request) {
   llvm::json::ObjectMapper O(Parameters);
   int64_t Limit;
   bool OK =
       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;
 }
 
 llvm::json::Value toJSON(const FuzzyFindRequest &Request) {
   return llvm::json::Object{
       {"Query", Request.Query},
       {"Scopes", llvm::json::Array{Request.Scopes}},
       {"AnyScope", Request.AnyScope},
       {"Limit", Request.Limit},
       {"RestrictForCodeCompletion", Request.RestrictForCodeCompletion},
       {"ProximityPaths", llvm::json::Array{Request.ProximityPaths}},
+      {"PreferredTypes", llvm::json::Array{Request.PreferredTypes}},
   };
 }
 
 bool SwapIndex::fuzzyFind(const FuzzyFindRequest &R,
                           llvm::function_ref<void(const Symbol &)> CB) const {
   return snapshot()->fuzzyFind(R, CB);
 }
 void SwapIndex::lookup(const LookupRequest &R,

clangd/index/dex/Dex.h

public:
   void refs(const RefsRequest &Req,
             llvm::function_ref<void(const Ref &)> Callback) const override;
 
   size_t estimateMemoryUsage() const override;
 
 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;
   /// SymbolQuality[I] is the quality of Symbols[I].
   std::vector<float> SymbolQuality;
   llvm::DenseMap<SymbolID, const Symbol *> LookupTable;
   /// Inverted index is a mapping from the search token to the posting list,
   /// which contains all items which can be characterized by such search token.

clangd/index/dex/Dex.cpp

 
 // Returns the tokens which are given symbol's characteristics. Currently, the
 // generated tokens only contain fuzzy matching trigrams and symbol's scope,
 // but in the future this will also return path proximity tokens and other
 // types of tokens such as symbol type (if applicable).
 // 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);
   Result.emplace_back(Token::Kind::Scope, Sym.Scope);
   // Skip token generation for symbols with unknown declaration location.
   if (!llvm::StringRef(Sym.CanonicalDeclaration.FileURI).empty())
     for (const auto &ProximityURI :
          generateProximityURIs(Sym.CanonicalDeclaration.FileURI))
       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() {
   this->Corpus = dex::Corpus(Symbols.size());
   std::vector<std::pair<float, const Symbol *>> ScoredSymbols(Symbols.size());
 
   for (size_t I = 0; I < Symbols.size(); ++I) {
     const Symbol *Sym = Symbols[I];
     LookupTable[Sym->ID] = Sym;
     ScoredSymbols[I] = {quality(*Sym), Sym};
   }
 
   // Symbols are sorted by symbol qualities so that items in the posting lists
   // are stored in the descending order of symbol quality.
   llvm::sort(ScoredSymbols, std::greater<std::pair<float, const Symbol *>>());
 
   // SymbolQuality was empty up until now.
   SymbolQuality.resize(Symbols.size());
   // Populate internal storage using Symbol + Score pairs.
   for (size_t I = 0; I < ScoredSymbols.size(); ++I) {

 }
 
 std::unique_ptr<Iterator> Dex::iterator(const Token &Tok) const {
   auto It = InvertedIndex.find(Tok);
   return It == InvertedIndex.end() ? Corpus.none()
                                    : 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.
 bool Dex::fuzzyFind(const FuzzyFindRequest &Req,
                     llvm::function_ref<void(const Symbol &)> Callback) const {
   assert(!StringRef(Req.Query).contains("::") &&
          "There must be no :: in query.");
   trace::Span Tracer("Dex fuzzyFind");
bool Dex::fuzzyFind(const FuzzyFindRequest &Req,
   for (const auto &Scope : Req.Scopes)
     ScopeIterators.push_back(iterator(Token(Token::Kind::Scope, Scope)));
   if (Req.AnyScope)
     ScopeIterators.push_back(
         Corpus.boost(Corpus.all(), ScopeIterators.empty() ? 1.0 : 0.2));
   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));
 
   // Use TRUE iterator if both trigrams and scopes from the query are not
   // present in the symbol index.
   auto Root = Corpus.intersect(move(Criteria));
   // Retrieve more items than it was requested: some of  the items with high

clangd/index/dex/Token.h

enum class Kind {
     Scope,
     /// Path Proximity URI to symbol declaration.
     ///
     /// Data stores path URI of symbol declaration file or its parent.
     ///
     /// 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)
       : Data(Data), TokenKind(TokenKind) {}
 
   bool operator==(const Token &Other) const {
     return TokenKind == Other.TokenKind && Data == Other.Data;
   }
 
   /// Representation which is unique among Token with the same Kind.
   std::string Data;
   Kind TokenKind;
 
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Token &T) {
     switch (T.TokenKind) {
     case Kind::Trigram:
       OS << "T=";
       break;
     case Kind::Scope:
       OS << "S=";
       break;
     case Kind::ProximityURI:
       OS << "U=";
       break;
+    case Kind::Type:
+      OS << "Ty=";
+      break;
     case Kind::Sentinel:
       OS << "?=";
       break;
     }
     return OS << T.Data;
   }
 
 private:

unittests/clangd/DexTests.cpp

TEST(DexTests, Refs) {
   Req.Limit = 1;
   Files.clear();
   Dex(std::vector<Symbol>{Foo, Bar}, Refs).refs(Req, [&](const Ref &R) {
     Files.push_back(R.Location.FileURI);
   });
   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
 } // namespace clang