[clangd] Don't create as much garbage while building Dex index.

[sammccall]

Summary

The Token objects are relatively expensive and we were spending a lot of
CPU creating them for each trigram emitted. Instead, use a tiny trigram
structure until we're ready to finalize the index.

This improves the new BuildDex benchmark by 20%. This code is hot and on
the critical path in clangd: it runs after a new preamble is built.

[kbobyrev]

Wow, nice work, thank you!

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

51

No longer "returns": "adds/generates"?

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

84

It isn't immediately obvious why processing different id lengths is different. I'm assuming it's less expensive to find duplicates in the small vector as we go than to find the unique elements in the end in case of the short identifier, but I'd be happy to have a comment here.

Also, 14 is a magic number:

• Move it to Trigram::SMALL_IDENTIFIER_SIZE or some other constant?
• How did you find this number? Did you try running benchmark with different values and figured this one is the best for the LLVM index? If not, it would make sense to do that.

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

35–55

Also doesn't "return" anymore

49

nit: maybe call it UID or something similar, "unsigned int" is somewhat surprising in this context

[sammccall]

Address comments, add more docs.

[sammccall]

Use a faster hash function, which seems to be worth something in the neighbourhood of 4%.

[kbobyrev]

LGTM, thank you for the patch!

clang-tools-extra/clangd/benchmarks/IndexBenchmark.cpp

static void DexQueries(benchmark::State &State) {
   const auto Dex = buildDex();
   const auto Requests = extractQueriesFromLogs();
   for (auto _ : State)
     for (const auto &Request : Requests)
       Dex->fuzzyFind(Request, [](const Symbol &S) {});
 }
 BENCHMARK(DexQueries);
 
+static void DexBuild(benchmark::State &State) {
+  for (auto _ : State)
+    buildDex();
+}
+BENCHMARK(DexBuild);
+
 } // namespace
 } // namespace clangd
 } // namespace clang
 
 // FIXME(kbobyrev): Add index building time benchmarks.
 // FIXME(kbobyrev): Add memory consumption "benchmarks" by manually measuring
 // in-memory index size and reporting it as time.
 // FIXME(kbobyrev): Create a logger wrapper to suppress debugging info printer.

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

 }
 
 namespace {
 
 // Mark symbols which are can be used for code completion.
 const Token RestrictedForCodeCompletion =
     Token(Token::Kind::Sentinel, "Restricted For Code Completion");
 
-// 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.
+// Helper to efficiently assemble the inverse index (token -> matching docs).
+// The output is a nice uniform structure keyed on Token, but constructing
+// the Token object every time we want to insert into the map is wasteful.
+// Instead we have various maps keyed on things that are cheap to compute,
+// and produce the Token keys once at the end.
+class IndexBuilder {
+  llvm::DenseMap<Trigram, std::vector<DocID>> TrigramDocs;
+  std::vector<DocID> RestrictedCCDocs;
+  llvm::StringMap<std::vector<DocID>> TypeDocs;
+  llvm::StringMap<std::vector<DocID>> ScopeDocs;
+  llvm::StringMap<std::vector<DocID>> ProximityDocs;
+  std::vector<Trigram> TrigramScratch;
+
+public:
+  // Add the tokens which are given symbol's characteristics.
+  // This includes fuzzy matching trigrams, symbol's scope, etc.
   // FIXME(kbobyrev): Support more token 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.
+  void add(const Symbol &Sym, DocID D) {
+    generateIdentifierTrigrams(Sym.Name, TrigramScratch);
+    for (Trigram T : TrigramScratch)
+      TrigramDocs[T].push_back(D);
+    ScopeDocs[Sym.Scope].push_back(D);
     if (!llvm::StringRef(Sym.CanonicalDeclaration.FileURI).empty())
       for (const auto &ProximityURI :
            generateProximityURIs(Sym.CanonicalDeclaration.FileURI))
-      Result.emplace_back(Token::Kind::ProximityURI, ProximityURI);
+        ProximityDocs[ProximityURI].push_back(D);
     if (Sym.Flags & Symbol::IndexedForCodeCompletion)
-    Result.emplace_back(RestrictedForCodeCompletion);
+      RestrictedCCDocs.push_back(D);
     if (!Sym.Type.empty())
-    Result.emplace_back(Token::Kind::Type, Sym.Type);
+      TypeDocs[Sym.Type].push_back(D);
+  }
+
+  // Assemble the final compressed posting lists for the added symbols.
+  llvm::DenseMap<Token, PostingList> build() {
+    llvm::DenseMap<Token, PostingList> Result(/*InitialReserve=*/
+                                              TrigramDocs.size() +
+                                              RestrictedCCDocs.size() +
+                                              TypeDocs.size() +
+                                              ScopeDocs.size() +
+                                              ProximityDocs.size());
+    for (const auto &E : TrigramDocs)
+      Result.try_emplace(Token(Token::Kind::Trigram, E.first.str()), E.second);
+    for (const auto &E : TypeDocs)
+      Result.try_emplace(Token(Token::Kind::Type, E.first()), E.second);
+    for (const auto &E : ScopeDocs)
+      Result.try_emplace(Token(Token::Kind::Scope, E.first()), E.second);
+    for (const auto &E : ProximityDocs)
+      Result.try_emplace(Token(Token::Kind::ProximityURI, E.first()), E.second);
+    if (!RestrictedCCDocs.empty())
+      Result.try_emplace(RestrictedForCodeCompletion, RestrictedCCDocs);
     return Result;
   }
+};
 
 } // 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) {
void Dex::buildIndex() {
   // 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) {
     SymbolQuality[I] = ScoredSymbols[I].first;
     Symbols[I] = ScoredSymbols[I].second;
   }
 
-  // Populate TempInvertedIndex with lists for index symbols.
-  llvm::DenseMap<Token, std::vector<DocID>> TempInvertedIndex;
-  for (DocID SymbolRank = 0; SymbolRank < Symbols.size(); ++SymbolRank) {
-    const auto *Sym = Symbols[SymbolRank];
-    for (const auto &Token : generateSearchTokens(*Sym))
-      TempInvertedIndex[Token].push_back(SymbolRank);
-  }
-
-  // Convert lists of items to posting lists.
-  for (const auto &TokenToPostingList : TempInvertedIndex)
-    InvertedIndex.insert(
-        {TokenToPostingList.first, PostingList(TokenToPostingList.second)});
+  // Build posting lists for symbols.
+  IndexBuilder Builder;
+  for (DocID SymbolRank = 0; SymbolRank < Symbols.size(); ++SymbolRank)
+    Builder.add(*Symbols[SymbolRank], SymbolRank);
+  InvertedIndex = Builder.build();
 }
 
 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);
 }
 

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

 /// identifier).
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_DEX_TRIGRAM_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DEX_TRIGRAM_H
 
 #include "Token.h"
+#include "llvm/ADT/bit.h"
 
 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace dex {
 
-/// Returns list of unique fuzzy-search trigrams from unqualified symbol.
+// Compact representation of a trigram (string with up to 3 characters).
+// Trigram generation is the hot path of indexing, so Token is too wasteful.
+class Trigram {
+  std::array<char, 4> Data; // Last element is length.
+  // Steal some invalid bit patterns for DenseMap sentinels.
+  enum class Sentinel { Tombstone = 4, Empty = 5 };
+  Trigram(Sentinel S) : Data{0, 0, 0, static_cast<char>(S)} {}
+  uint32_t id() const { return llvm::bit_cast<uint32_t>(Data); }
+
+public:
+  Trigram() : Data{0, 0, 0, 0} {}
+  Trigram(char A) : Data{A, 0, 0, 1} {}
+  Trigram(char A, char B) : Data{A, B, 0, 2} {}
+  Trigram(char A, char B, char C) : Data{A, B, C, 3} {}
+  std::string str() const { return std::string(Data.data(), Data[3]); }
+  friend struct ::llvm::DenseMapInfo<Trigram>;
+  friend bool operator==(Trigram L, Trigram R) { return L.id() == R.id(); }
+  friend bool operator<(Trigram L, Trigram R) { return L.id() < R.id(); }
+};
+
+/// Produces list of unique fuzzy-search trigrams from unqualified symbol.
 /// The trigrams give the 3-character query substrings this symbol can match.
 ///
 /// The symbol's name is broken into segments, e.g. "FooBar" has two segments.
 /// Trigrams can start at any character in the input. Then we can choose to move
 /// to the next character or to the start of the next segment.
 ///
 /// Short trigrams (length 1-2) are used for short queries. These are:
 ///  - prefixes of the identifier, of length 1 and 2
 ///  - the first character + next head character
 ///
 /// For "FooBar" we get the following trigrams:
 ///  {f, fo, fb, foo, fob, fba, oob, oba, bar}.
 ///
 /// Trigrams are lowercase, as trigram matching is case-insensitive.
-/// Trigrams in the returned list are deduplicated.
-std::vector<Token> generateIdentifierTrigrams(llvm::StringRef Identifier);
+/// Trigrams in the list are deduplicated.
+void generateIdentifierTrigrams(llvm::StringRef Identifier,
+                                std::vector<Trigram> &Out);
 
 /// Returns list of unique fuzzy-search trigrams given a query.
 ///
 /// Query is segmented using FuzzyMatch API and downcasted to lowercase. Then,
 /// the simplest trigrams - sequences of three consecutive letters and digits
 /// are extracted and returned after deduplication.
 ///
 /// For short queries (less than 3 characters with Head or Tail roles in Fuzzy
 /// Matching segmentation) this returns a single trigram with the first
 /// characters (up to 3) to perform prefix match.
 std::vector<Token> generateQueryTrigrams(llvm::StringRef Query);
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang
 
+namespace llvm {
+template <> struct llvm::DenseMapInfo<clang::clangd::dex::Trigram> {
+  using Trigram = clang::clangd::dex::Trigram;
+  static inline Trigram getEmptyKey() {
+    return Trigram(Trigram::Sentinel::Empty);
+  }
+  static inline Trigram getTombstoneKey() {
+    return Trigram(Trigram::Sentinel::Tombstone);
+  }
+  static unsigned getHashValue(Trigram V) {
+    // Finalize step from MurmurHash3.
+    uint32_t X = V.id();
+    X ^= X >> 16;
+    X *= uint32_t{0x85ebca6b};
+    X ^= X >> 13;
+    X *= uint32_t{0xc2b2ae35};
+    X ^= X >> 16;
+    return X;
+  }
+  static bool isEqual(const Trigram &LHS, const Trigram &RHS) {
+    return LHS == RHS;
+  }
+};
+} // namespace llvm
+
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_DEX_TRIGRAM_H

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/StringExtras.h"
 #include <cctype>
 #include <queue>
 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace dex {
 
-std::vector<Token> generateIdentifierTrigrams(llvm::StringRef Identifier) {
+// 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());
   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
static void identifierTrigrams(llvm::StringRef Identifier, Func Out) {
   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;
     }
   }
 
-  llvm::DenseSet<Token> UniqueTrigrams;
-
-  auto Add = [&](std::string Chars) {
-    UniqueTrigrams.insert(Token(Token::Kind::Trigram, Chars));
-  };
-
   // Iterate through valid sequences of three characters Fuzzy Matcher can
   // process.
   for (size_t I = 0; I < LowercaseIdentifier.size(); ++I) {
     // Skip delimiters.
     if (Roles[I] != Head && Roles[I] != Tail)
       continue;
     for (const unsigned J : Next[I]) {
       if (J == 0)
         continue;
       for (const unsigned K : Next[J]) {
         if (K == 0)
           continue;
-        Add({{LowercaseIdentifier[I], LowercaseIdentifier[J],
-              LowercaseIdentifier[K]}});
+        Out(Trigram(LowercaseIdentifier[I], LowercaseIdentifier[J],
+                    LowercaseIdentifier[K]));
       }
     }
   }
   // Emit short-query trigrams: FooBar -> f, fo, fb.
   if (!LowercaseIdentifier.empty())
-    Add({LowercaseIdentifier[0]});
+    Out(Trigram(LowercaseIdentifier[0]));
   if (LowercaseIdentifier.size() >= 2)
-    Add({LowercaseIdentifier[0], LowercaseIdentifier[1]});
+    Out(Trigram(LowercaseIdentifier[0], LowercaseIdentifier[1]));
   for (size_t I = 1; I < LowercaseIdentifier.size(); ++I)
     if (Roles[I] == Head) {
-      Add({LowercaseIdentifier[0], LowercaseIdentifier[I]});
+      Out(Trigram(LowercaseIdentifier[0], LowercaseIdentifier[I]));
       break;
     }
+}
 
-  return {UniqueTrigrams.begin(), UniqueTrigrams.end()};
+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)};

clang-tools-extra/clangd/test/Inputs/requests.json

-[{"Limit":100,"ProximityPaths":["/usr/home/user/clang-tools-extra/clangd/benchmarks/IndexBenchmark.cpp"],"Query":"OMP","RestrictForCodeCompletion":true,"Scopes":["clang::"], "AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"s","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"sy","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"sys","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"sys","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"Dex","RestrictForCodeCompletion":true,"Scopes":["clang::clangd::", "clang::", "clang::clangd::dex::"],"AnyScope":false},
-{"Limit":100,"ProximityPaths":[],"Query":"Variable","RestrictForCodeCompletion":true,"Scopes":[""], "AnyScope":false}]
+[{"Limit":100,"ProximityPaths":["/usr/home/user/clang-tools-extra/clangd/benchmarks/IndexBenchmark.cpp"],"Query":"OMP","RestrictForCodeCompletion":true,"Scopes":["clang::"], "AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"s","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"sy","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"sys","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"sys","RestrictForCodeCompletion":true,"Scopes":["llvm::", ""], "AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"Dex","RestrictForCodeCompletion":true,"Scopes":["clang::clangd::", "clang::", "clang::clangd::dex::"],"AnyScope":false, "PreferredTypes":[]},
+{"Limit":100,"ProximityPaths":[],"Query":"Variable","RestrictForCodeCompletion":true,"Scopes":[""], "AnyScope":false, "PreferredTypes":[]}]

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

tokensAre(std::initializer_list<std::string> Strings, Token::Kind Kind) {
   return ::testing::UnorderedElementsAreArray(Tokens);
 }
 
 ::testing::Matcher<std::vector<Token>>
 trigramsAre(std::initializer_list<std::string> Trigrams) {
   return tokensAre(Trigrams, Token::Kind::Trigram);
 }
 
+std::vector<Token> identifierTrigramTokens(llvm::StringRef S) {
+  std::vector<Trigram> Trigrams;
+  generateIdentifierTrigrams(S, Trigrams);
+  std::vector<Token> Tokens;
+  for (Trigram T : Trigrams)
+    Tokens.emplace_back(Token::Kind::Trigram, T.str());
+  return Tokens;
+}
+
 TEST(DexTrigrams, IdentifierTrigrams) {
-  EXPECT_THAT(generateIdentifierTrigrams("X86"),
-              trigramsAre({"x86", "x", "x8"}));
+  EXPECT_THAT(identifierTrigramTokens("X86"), trigramsAre({"x86", "x", "x8"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("nl"), trigramsAre({"nl", "n"}));
+  EXPECT_THAT(identifierTrigramTokens("nl"), trigramsAre({"nl", "n"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("n"), trigramsAre({"n"}));
+  EXPECT_THAT(identifierTrigramTokens("n"), trigramsAre({"n"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("clangd"),
+  EXPECT_THAT(identifierTrigramTokens("clangd"),
               trigramsAre({"c", "cl", "cla", "lan", "ang", "ngd"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("abc_def"),
+  EXPECT_THAT(identifierTrigramTokens("abc_def"),
               trigramsAre({"a", "ab", "ad", "abc", "abd", "ade", "bcd", "bde",
                            "cde", "def"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("a_b_c_d_e_"),
+  EXPECT_THAT(identifierTrigramTokens("a_b_c_d_e_"),
               trigramsAre({"a", "a_", "ab", "abc", "bcd", "cde"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("unique_ptr"),
+  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(
-      generateIdentifierTrigrams("TUDecl"),
+      identifierTrigramTokens("TUDecl"),
       trigramsAre({"t", "tu", "td", "tud", "tde", "ude", "dec", "ecl"}));
 
-  EXPECT_THAT(generateIdentifierTrigrams("IsOK"),
+  EXPECT_THAT(identifierTrigramTokens("IsOK"),
               trigramsAre({"i", "is", "io", "iso", "iok", "sok"}));
 
   EXPECT_THAT(
-      generateIdentifierTrigrams("abc_defGhij__klm"),
+      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"}));
 }
 
 TEST(DexTrigrams, QueryTrigrams) {
   EXPECT_THAT(generateQueryTrigrams("c"), trigramsAre({"c"}));