[clangd] Dex Trigrams: Improve query trigram generation

[kbobyrev]

Summary

These are the trigrams for queries right now:

• "va" -> {Trigram("va")}
• "va_" -> {} (empty)

This is suboptimal since the resulting query will discard the query information
and return all symbols, some of which will be later be scored expensively
(fuzzy matching score). This is related to
https://github.com/clangd/clangd/issues/39 but does not fix it. Accidentally,
because of that incorrect behavior, when user types "tok::va" there are no
results (the issue is that tok::kw___builtin_va_arg does not have "va" token)
but when "tok::va_" is typed, expected result (tok::kw___builtin_va_arg)
shows up by accident. This is because the dex query transformer will only
lookup symbols within the tok:: namespace. There won't be many, so the
returned results will contain symbol we need; this symbol will be filtered out
by the expensive checks and that will be displayed in the editor.

[kbobyrev]

Add a small comment.

[kbobyrev]

Add improvements for identifier trigram generation and make query & id
generators consistent with each other.

[sammccall]

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

73–93

It'd be good to describe what the semantics actually are somewhere.
As it is, this comment describes the algorithm used to generate the trigrams, and the other function refers to this one.

76

This seems overly complicated.
If a separator precedes the alphanums, then it's at the front.
Moreover its Next table seems to be set up correctly.

So I believe this all amounts to: if the first char is a separator, then we allow starting our partial trigrams at 0, first head, second head. Else just first head & second head. And we just follow the next table.

100

This looks a lot like the Next table. Can't we use that? (Second head + third head is missing, but I don't see why)

128

both "valid" and "symbols" are confusing here.

Do we actually need this variable?
Seems like we can detect the case at the end, if we generated no trigrams.

133

I'd suggest always including the separator, and then using StringRef(Letters).take_back(2). Then we don't need ValidSymbols here either.

[kbobyrev]

Address review comments.

[kbobyrev]

Slightly improve docs.

[sammccall]

Basically LG now, some of the changes are a bit mysterious to me though

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

49

why the change from array to pair and 0 to NPOS?

103

Fist -> first
You're missing first separator only

I think examples might be clearer than a description of these.

123

(if the change to NPOS is just to avoid hitting Position == 0 on the first iteration, you could check it here instead)

140

Since the query trigram corresponds so closely to the query for short queries, I think it makes more sense to consider it the other way: generateIdentifierTrigrams is deciding which queries it wants to match and emulating this function.

143

nit: by prepending a separator, we form a bigram!
If a bigram can't be formed out of letters/numbers...

[kbobyrev]

Address review comments.

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

49

• Array of two elements is confusing: we had it since we had 3 elements, with 2 it doesn't make sense anymore. Also, makes it easier to decompose the items (std::tie(Tail, Head) vs size_t Tail = Next[I][0], Head = Next[I][1]).
• NPOS is more explicit than 0 in our intent: 0 is a valid index and it feels awkward to use a valid index as NOT_FOUND. Yes, it's the one that can't appear in the Next but it's not clear from the code and requires some explanation for understanding. Why NPOS can't be found in Next is obvious.

123

It's not only for that but this was the trigger, having the Position == 0 check here makes the logic slightly confusing, I think it's better to just use NPOS instead for the reasons above.

[sammccall]

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

49

I disagree on the style points here and they don't seem particularly related to the patch. Do they need to be included here?

Array of two elements is confusing

Well, I disagree here: the difference between array<T, 2> and pair<T, T> is homogeneous vs heterogeneous. The usage here is almost entirely homogeneous.
3 vs 2 doesn't make a difference, we could have used tuple<T, T, T>, same would apply.

Also, makes it easier to decompose the items

As far as ergonomics go, this comes at the cost of not being able to iterate over them, which we do much more often. (And AFAICT you never need NextTail decomposed.)

NPOS is not a common convention so I don't think it's more obvious.

49

you've changed unsigned -> size_t here, can you revert that please?

112

You're not treating head vs tail differently, this should just be iterating over Next as above.

[kbobyrev]

Address review comments.

[kbobyrev]

Remove refactoring leftovers.

[sammccall]

Sorry about the delay!

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

46

the default size of this vector is 6 - if we care about the optimization we probably want to double it or so?

(Roles above is fine, it's 40)

[kbobyrev]

No problem, thank you!

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

 //===--- Trigram.cpp - Trigram generation for Fuzzy Matching ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Trigram.h"
 #include "FuzzyMatch.h"
 #include "Token.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include <cctype>
+#include <limits>
 #include <queue>
 #include <string>
+#include <utility>
 
 namespace clang {
 namespace clangd {
 namespace dex {
 
 // Produce trigrams (including duplicates) and pass them to Out().
 template <typename Func>
 static void identifierTrigrams(llvm::StringRef Identifier, Func Out) {
   // Apply fuzzy matching text segmentation.
-  std::vector<CharRole> Roles(Identifier.size());
+  llvm::SmallVector<CharRole> Roles(Identifier.size());
   calculateRoles(Identifier,
                  llvm::makeMutableArrayRef(Roles.data(), Identifier.size()));
 
   std::string LowercaseIdentifier = Identifier.lower();
 
   // For each character, store indices of the characters to which fuzzy matching
-  // algorithm can jump. There are 3 possible variants:
+  // algorithm can jump. There are 2 possible variants:
   //
   // * Next Tail - next character from the same segment
   // * Next Head - front character of the next segment
   //
   // Next stores tuples of three indices in the presented order, if a variant is
   // not available then 0 is stored.
-  std::vector<std::array<unsigned, 3>> Next(LowercaseIdentifier.size());
+  llvm::SmallVector<std::array<unsigned, 2>> Next(LowercaseIdentifier.size());
   unsigned NextTail = 0, NextHead = 0;
   for (int I = LowercaseIdentifier.size() - 1; I >= 0; --I) {
     Next[I] = {{NextTail, NextHead}};
     NextTail = Roles[I] == Tail ? I : 0;
     if (Roles[I] == Head) {
       NextHead = I;
     }
   }
 
   // Iterate through valid sequences of three characters Fuzzy Matcher can
   // process.
-  for (size_t I = 0; I < LowercaseIdentifier.size(); ++I) {
+  for (unsigned I = 0; I < LowercaseIdentifier.size(); ++I) {
     // Skip delimiters.
     if (Roles[I] != Head && Roles[I] != Tail)
       continue;
-    for (const unsigned J : Next[I]) {
+    for (unsigned J : Next[I]) {
       if (J == 0)
         continue;
-      for (const unsigned K : Next[J]) {
+      for (unsigned K : Next[J]) {
         if (K == 0)
           continue;
         Out(Trigram(LowercaseIdentifier[I], LowercaseIdentifier[J],
                     LowercaseIdentifier[K]));
       }
     }
   }
-  // Emit short-query trigrams: FooBar -> f, fo, fb.
-  if (!LowercaseIdentifier.empty())
-    Out(Trigram(LowercaseIdentifier[0]));
-  if (LowercaseIdentifier.size() >= 2)
-    Out(Trigram(LowercaseIdentifier[0], LowercaseIdentifier[1]));
-  for (size_t I = 1; I < LowercaseIdentifier.size(); ++I)
-    if (Roles[I] == Head) {
-      Out(Trigram(LowercaseIdentifier[0], LowercaseIdentifier[I]));
+  // Short queries semantics are different. When the user dosn't type enough
+  // symbols to form trigrams, we still want to serve meaningful results. To
+  // achieve that, we form incomplete trigrams (bi- and unigrams) for the
+  // identifiers and also generate short trigrams on the query side from what
+  // is available. We allow a small number of short trigram types in order to
+  // prevent excessive memory usage and increase the quality of the search.
+  // Only the first few symbols are allowed to be used in incomplete trigrams.
+  //
+  // Example - for "_abc_def_ghi_jkl" we'll get following incomplete trigrams:
+  // "_", "_a", "a", "ab", "ad", "d", "de", "dg"
+  for (unsigned Position = 0, HeadsSeen = 0; HeadsSeen < 2;) {
+    // The first symbol might be a separator, so the loop condition should be
+    // stopping as soon as there is no next head or we have seen two heads.
+    if (Roles[Position] == Head)
+      ++HeadsSeen;
+    Out(Trigram(LowercaseIdentifier[Position]));
+    for (unsigned I : Next[Position])
+      if (I != 0)
+        Out(Trigram(LowercaseIdentifier[Position], LowercaseIdentifier[I]));
+    Position = Next[Position][1];
+    if (Position == 0)
       break;
-    }
 }
 
 void generateIdentifierTrigrams(llvm::StringRef Identifier,
                                 std::vector<Trigram> &Result) {
   // Empirically, scanning for duplicates is faster with fewer trigrams, and
   // a hashtable is faster with more. This is a hot path, so dispatch based on
   // expected number of trigrams. Longer identifiers produce more trigrams.
   // The magic number was tuned by running IndexBenchmark.DexBuild.
   constexpr unsigned ManyTrigramsIdentifierThreshold = 14;
   Result.clear();
   if (Identifier.size() < ManyTrigramsIdentifierThreshold) {
     identifierTrigrams(Identifier, [&](Trigram T) {
       if (!llvm::is_contained(Result, T))
         Result.push_back(T);
     });
   } else {
     identifierTrigrams(Identifier, [&](Trigram T) { Result.push_back(T); });
     llvm::sort(Result);
     Result.erase(std::unique(Result.begin(), Result.end()), Result.end());
   }
 }
 
 std::vector<Token> generateQueryTrigrams(llvm::StringRef Query) {
   if (Query.empty())
     return {};
-  std::string LowercaseQuery = Query.lower();
-  if (Query.size() < 3) // short-query trigrams only
-    return {Token(Token::Kind::Trigram, LowercaseQuery)};
 
   // Apply fuzzy matching text segmentation.
-  std::vector<CharRole> Roles(Query.size());
+  llvm::SmallVector<CharRole> Roles(Query.size());
   calculateRoles(Query, llvm::makeMutableArrayRef(Roles.data(), Query.size()));
 
+  std::string LowercaseQuery = Query.lower();
+
   llvm::DenseSet<Token> UniqueTrigrams;
   std::string Chars;
-  for (unsigned I = 0; I < Query.size(); ++I) {
+  for (unsigned I = 0; I < LowercaseQuery.size(); ++I) {
     if (Roles[I] != Head && Roles[I] != Tail)
       continue; // Skip delimiters.
     Chars.push_back(LowercaseQuery[I]);
     if (Chars.size() > 3)
       Chars.erase(Chars.begin());
     if (Chars.size() == 3)
       UniqueTrigrams.insert(Token(Token::Kind::Trigram, Chars));
   }
 
+  // For queries with very few letters, generateIdentifierTrigrams emulates
+  // outputs of this function to match the semantics.
+  if (UniqueTrigrams.empty()) {
+    // If bigram can't be formed out of letters/numbers, we prepend separator.
+    std::string Result(1, LowercaseQuery.front());
+    for (unsigned I = 1; I < LowercaseQuery.size(); ++I)
+      if (Roles[I] == Head || Roles[I] == Tail)
+        Result += LowercaseQuery[I];
+    UniqueTrigrams.insert(
+        Token(Token::Kind::Trigram, llvm::StringRef(Result).take_back(2)));
+  }
+
   return {UniqueTrigrams.begin(), UniqueTrigrams.end()};
 }
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang

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

TEST(DexTrigrams, IdentifierTrigrams) {
   EXPECT_THAT(identifierTrigramTokens("nl"), trigramsAre({"nl", "n"}));
 
   EXPECT_THAT(identifierTrigramTokens("n"), trigramsAre({"n"}));
 
   EXPECT_THAT(identifierTrigramTokens("clangd"),
               trigramsAre({"c", "cl", "cla", "lan", "ang", "ngd"}));
 
   EXPECT_THAT(identifierTrigramTokens("abc_def"),
-              trigramsAre({"a", "ab", "ad", "abc", "abd", "ade", "bcd", "bde",
-                           "cde", "def"}));
+              trigramsAre({"a", "d", "ab", "ad", "de", "abc", "abd", "ade",
+                           "bcd", "bde", "cde", "def"}));
 
   EXPECT_THAT(identifierTrigramTokens("a_b_c_d_e_"),
-              trigramsAre({"a", "a_", "ab", "abc", "bcd", "cde"}));
+              trigramsAre({"a", "b", "ab", "bc", "abc", "bcd", "cde"}));
 
   EXPECT_THAT(identifierTrigramTokens("unique_ptr"),
-              trigramsAre({"u", "un", "up", "uni", "unp", "upt", "niq", "nip",
-                           "npt", "iqu", "iqp", "ipt", "que", "qup", "qpt",
-                           "uep", "ept", "ptr"}));
-
-  EXPECT_THAT(
-      identifierTrigramTokens("TUDecl"),
-      trigramsAre({"t", "tu", "td", "tud", "tde", "ude", "dec", "ecl"}));
+              trigramsAre({"u",   "p",   "un",  "up",  "pt",  "uni", "unp",
+                           "upt", "niq", "nip", "npt", "iqu", "iqp", "ipt",
+                           "que", "qup", "qpt", "uep", "ept", "ptr"}));
+
+  EXPECT_THAT(identifierTrigramTokens("TUDecl"),
+              trigramsAre({"t", "d", "tu", "td", "de", "tud", "tde", "ude",
+                           "dec", "ecl"}));
 
   EXPECT_THAT(identifierTrigramTokens("IsOK"),
-              trigramsAre({"i", "is", "io", "iso", "iok", "sok"}));
+              trigramsAre({"i", "o", "is", "ok", "io", "iso", "iok", "sok"}));
 
-  EXPECT_THAT(
-      identifierTrigramTokens("abc_defGhij__klm"),
-      trigramsAre({"a",   "ab",  "ad",  "abc", "abd", "ade", "adg", "bcd",
-                   "bde", "bdg", "cde", "cdg", "def", "deg", "dgh", "dgk",
-                   "efg", "egh", "egk", "fgh", "fgk", "ghi", "ghk", "gkl",
-                   "hij", "hik", "hkl", "ijk", "ikl", "jkl", "klm"}));
+  EXPECT_THAT(identifierTrigramTokens("_pb"),
+              trigramsAre({"_", "_p", "p", "pb"}));
+  EXPECT_THAT(identifierTrigramTokens("__pb"),
+              trigramsAre({"_", "_p", "p", "pb"}));
+
+  EXPECT_THAT(identifierTrigramTokens("abc_defGhij__klm"),
+              trigramsAre({"a",   "d",   "ab",  "ad",  "dg",  "de",  "abc",
+                           "abd", "ade", "adg", "bcd", "bde", "bdg", "cde",
+                           "cdg", "def", "deg", "dgh", "dgk", "efg", "egh",
+                           "egk", "fgh", "fgk", "ghi", "ghk", "gkl", "hij",
+                           "hik", "hkl", "ijk", "ikl", "jkl", "klm"}));
 }
 
 TEST(DexTrigrams, QueryTrigrams) {
   EXPECT_THAT(generateQueryTrigrams("c"), trigramsAre({"c"}));
   EXPECT_THAT(generateQueryTrigrams("cl"), trigramsAre({"cl"}));
   EXPECT_THAT(generateQueryTrigrams("cla"), trigramsAre({"cla"}));
 
   EXPECT_THAT(generateQueryTrigrams(""), trigramsAre({}));
   EXPECT_THAT(generateQueryTrigrams("_"), trigramsAre({"_"}));
-  EXPECT_THAT(generateQueryTrigrams("__"), trigramsAre({"__"}));
-  EXPECT_THAT(generateQueryTrigrams("___"), trigramsAre({}));
+  EXPECT_THAT(generateQueryTrigrams("__"), trigramsAre({"_"}));
+  EXPECT_THAT(generateQueryTrigrams("___"), trigramsAre({"_"}));
+
+  EXPECT_THAT(generateQueryTrigrams("m_"), trigramsAre({"m"}));
+
+  EXPECT_THAT(generateQueryTrigrams("p_b"), trigramsAre({"pb"}));
+  EXPECT_THAT(generateQueryTrigrams("pb_"), trigramsAre({"pb"}));
+  EXPECT_THAT(generateQueryTrigrams("_p"), trigramsAre({"_p"}));
+  EXPECT_THAT(generateQueryTrigrams("_pb_"), trigramsAre({"pb"}));
+  EXPECT_THAT(generateQueryTrigrams("__pb"), trigramsAre({"pb"}));
 
   EXPECT_THAT(generateQueryTrigrams("X86"), trigramsAre({"x86"}));
 
   EXPECT_THAT(generateQueryTrigrams("clangd"),
               trigramsAre({"cla", "lan", "ang", "ngd"}));
 
   EXPECT_THAT(generateQueryTrigrams("abc_def"),
               trigramsAre({"abc", "bcd", "cde", "def"}));
TEST(DexTest, FuzzyMatch) {
   Req.Query = "lol";
   Req.AnyScope = true;
   Req.Limit = 2;
   EXPECT_THAT(match(*I, Req),
               UnorderedElementsAre("LaughingOutLoud", "LittleOldLady"));
 }
 
 TEST(DexTest, ShortQuery) {
-  auto I = Dex::build(generateSymbols({"OneTwoThreeFour"}), RefSlab(),
+  auto I = Dex::build(generateSymbols({"_OneTwoFourSix"}), RefSlab(),
                       RelationSlab());
   FuzzyFindRequest Req;
   Req.AnyScope = true;
   bool Incomplete;
 
-  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("OneTwoThreeFour"));
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
   EXPECT_FALSE(Incomplete) << "Empty string is not a short query";
 
-  Req.Query = "t";
-  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre());
-  EXPECT_TRUE(Incomplete) << "Short queries have different semantics";
+  Req.Query = "o";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using first head as unigram";
+
+  Req.Query = "_o";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using delimiter and first head as bigram";
+
+  Req.Query = "on";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using first head and tail as bigram";
+
+  Req.Query = "ot";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using first two heads as bigram";
+
+  Req.Query = "tw";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using second head and tail as bigram";
+
+  Req.Query = "tf";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
+  EXPECT_TRUE(Incomplete) << "Using second and third heads as bigram";
 
-  Req.Query = "tt";
+  Req.Query = "fo";
   EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre());
   EXPECT_TRUE(Incomplete) << "Short queries have different semantics";
 
-  Req.Query = "ttf";
-  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("OneTwoThreeFour"));
+  Req.Query = "tfs";
+  EXPECT_THAT(match(*I, Req, &Incomplete), ElementsAre("_OneTwoFourSix"));
   EXPECT_FALSE(Incomplete) << "3-char string is not a short query";
 }
 
 TEST(DexTest, MatchQualifiedNamesWithoutSpecificScope) {
   auto I = Dex::build(generateSymbols({"a::y1", "b::y2", "y3"}), RefSlab(),
                       RelationSlab());
   FuzzyFindRequest Req;
   Req.AnyScope = true;