[clangd] Avoid wasteful data structures in RefSlab::Builder

[sammccall]

Summary

This is worth another 10% or so on InedxBenchmark.DexBuild.

[kbobyrev]

LGTM with a couple of nits.

Thank you for the patch, this is a good idea!

clang-tools-extra/clangd/index/Ref.cpp

11

This defines BumpPtrAllocator, right? I think it should be included in clang-tools-extra/clangd/index/Ref.h, not here if the intent is to explicitly include used headers. Same for StringSaver.h, this source file does not use any additional variables of these types.

clang-tools-extra/clangd/index/Ref.h

141

nit: I know you like the short names, but with R I had to get back to Entry definition each time I saw it (I was thinking "wait, R? is there an L somewhere here?". Maybe it's just me, but one-letter fields seem scary. Maybe Reference?

[sammccall]

clang-tools-extra/clangd/index/Ref.cpp

11

Yup, that's right, thanks.
(An intermediate version of the code used the overloaded new from allocator.h)

clang-tools-extra/clangd/index/Ref.h

141

Oops, I agree but forgot to do this, thanks!

clang-tools-extra/clangd/index/Ref.h

 //===--- Ref.h ---------------------------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REF_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REF_H
 
 #include "SymbolID.h"
 #include "SymbolLocation.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 #include <set>
 #include <utility>
 
 namespace clang {
 namespace clangd {
public:
   public:
     Builder() : UniqueStrings(Arena) {}
     /// Adds a ref to the slab. Deep copy: Strings will be owned by the slab.
     void insert(const SymbolID &ID, const Ref &S);
     /// Consumes the builder to finalize the slab.
     RefSlab build() &&;
 
   private:
+    // A ref we're storing with its symbol to consume with build().
+    // All strings are interned, so DenseMapInfo can use pointer comparisons.
+    struct Entry {
+      SymbolID Symbol;
+      Ref Reference;
+    };
+    friend struct llvm::DenseMapInfo<Entry>;
+
     llvm::BumpPtrAllocator Arena;
     llvm::UniqueStringSaver UniqueStrings; // Contents on the arena.
-    llvm::DenseMap<SymbolID, std::set<Ref>> Refs;
+    llvm::DenseSet<Entry> Entries;
   };
 
 private:
   RefSlab(std::vector<value_type> Refs, llvm::BumpPtrAllocator Arena,
           size_t NumRefs)
       : Arena(std::move(Arena)), Refs(std::move(Refs)), NumRefs(NumRefs) {}
 
   llvm::BumpPtrAllocator Arena;
   std::vector<value_type> Refs;
   /// Number of all references.
   size_t NumRefs = 0;
 };
 
 } // namespace clangd
 } // namespace clang
 
+namespace llvm {
+template <> struct DenseMapInfo<clang::clangd::RefSlab::Builder::Entry> {
+  using Entry = clang::clangd::RefSlab::Builder::Entry;
+  static inline Entry getEmptyKey() {
+    static Entry E{clang::clangd::SymbolID(""), {}};
+    return E;
+  }
+  static inline Entry getTombstoneKey() {
+    static Entry E{clang::clangd::SymbolID("TOMBSTONE"), {}};
+    return E;
+  }
+  static unsigned getHashValue(const Entry &Val) {
+    return llvm::hash_combine(
+        Val.Symbol, reinterpret_cast<uintptr_t>(Val.Reference.Location.FileURI),
+        Val.Reference.Location.Start.rep(), Val.Reference.Location.End.rep());
+  }
+  static bool isEqual(const Entry &LHS, const Entry &RHS) {
+    return std::tie(LHS.Symbol, LHS.Reference.Location.FileURI,
+                    LHS.Reference.Kind) ==
+               std::tie(RHS.Symbol, RHS.Reference.Location.FileURI,
+                        RHS.Reference.Kind) &&
+           LHS.Reference.Location.Start == RHS.Reference.Location.Start &&
+           LHS.Reference.Location.End == RHS.Reference.Location.End;
+  }
+};
+}; // namespace llvm
+
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REF_H

clang-tools-extra/clangd/index/Ref.cpp

 //===--- Ref.cpp -------------------------------------------------*- 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 "Ref.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace clang {
 namespace clangd {
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, RefKind K) {
   if (K == RefKind::Unknown)
     return OS << "Unknown";
   static constexpr std::array<const char *, 4> Messages = {"Decl", "Def", "Ref",
                                                            "Spelled"};
llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, RefKind K) {
   return OS;
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Ref &R) {
   return OS << R.Location << ":" << R.Kind;
 }
 
 void RefSlab::Builder::insert(const SymbolID &ID, const Ref &S) {
-  auto &M = Refs[ID];
-  if (M.count(S))
-    return;
-  Ref R = S;
-  R.Location.FileURI =
-      UniqueStrings.save(R.Location.FileURI).data();
-  M.insert(std::move(R));
+  Entry E = {ID, S};
+  E.Reference.Location.FileURI = UniqueStrings.save(S.Location.FileURI).data();
+  Entries.insert(std::move(E));
 }
 
 RefSlab RefSlab::Builder::build() && {
-  // We can reuse the arena, as it only has unique strings and we need them all.
-  // Reallocate refs on the arena to reduce waste and indirections when reading.
   std::vector<std::pair<SymbolID, llvm::ArrayRef<Ref>>> Result;
-  Result.reserve(Refs.size());
-  size_t NumRefs = 0;
-  for (auto &Sym : Refs) {
-    std::vector<Ref> SymRefs(Sym.second.begin(), Sym.second.end());
-    NumRefs += SymRefs.size();
-    Result.emplace_back(Sym.first, llvm::ArrayRef<Ref>(SymRefs).copy(Arena));
+  // We'll reuse the arena, as it only has unique strings and we need them all.
+  // We need to group refs by symbol and form contiguous arrays on the arena.
+  std::vector<std::pair<SymbolID, const Ref *>> Flat;
+  Flat.reserve(Entries.size());
+  for (const Entry &E : Entries)
+    Flat.emplace_back(E.Symbol, &E.Reference);
+  // Group by SymbolID.
+  llvm::sort(Flat, llvm::less_first());
+  std::vector<Ref> Refs;
+  // Loop over symbols, copying refs for each onto the arena.
+  for (auto I = Flat.begin(), End = Flat.end(); I != End;) {
+    SymbolID Sym = I->first;
+    Refs.clear();
+    do {
+      Refs.push_back(*I->second);
+      ++I;
+    } while (I != End && I->first == Sym);
+    llvm::sort(Refs); // By file, affects xrefs display order.
+    Result.emplace_back(Sym, llvm::ArrayRef<Ref>(Refs).copy(Arena));
   }
-  return RefSlab(std::move(Result), std::move(Arena), NumRefs);
+  return RefSlab(std::move(Result), std::move(Arena), Entries.size());
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/SymbolLocation.h

struct SymbolLocation {
   //  - Column is (by default) the number of UTF-16 code between the last \n
   //    (or start of file) and the point.
   //    If the `offsetEncoding` protocol extension is used to negotiate UTF-8,
   //    then it is instead the number of *bytes* since the last \n.
   //
   // Position is encoded into 32 bits to save space.
   // If Line/Column overflow, the value will be their maximum value.
   struct Position {
-    Position() : Line(0), Column(0) {}
+    Position() : LineColumnPacked(0) {}
     void setLine(uint32_t Line);
-    uint32_t line() const { return Line; }
+    uint32_t line() const { return LineColumnPacked >> ColumnBits; }
     void setColumn(uint32_t Column);
-    uint32_t column() const { return Column; }
+    uint32_t column() const { return LineColumnPacked & MaxColumn; }
+    uint32_t rep() const { return LineColumnPacked; }
 
     bool hasOverflow() const {
-      return Line >= MaxLine || Column >= MaxColumn;
+      return line() == MaxLine || column() == MaxColumn;
     }
 
-    static constexpr uint32_t MaxLine = (1 << 20) - 1;
-    static constexpr uint32_t MaxColumn = (1 << 12) - 1;
+    static constexpr unsigned ColumnBits = 12;
+    static constexpr uint32_t MaxLine = (1 << (32 - ColumnBits)) - 1;
+    static constexpr uint32_t MaxColumn = (1 << ColumnBits) - 1;
 
   private:
-    uint32_t Line : 20;   // 0-based
-    uint32_t Column : 12; // 0-based
+    uint32_t LineColumnPacked; // Top 20 bit line, bottom 12 bits column.
   };
 
   /// The symbol range, using half-open range [Start, End).
   Position Start;
   Position End;
 
   explicit operator bool() const { return !llvm::StringRef(FileURI).empty(); }
 

clang-tools-extra/clangd/index/SymbolLocation.cpp

 
 namespace clang {
 namespace clangd {
 
 constexpr uint32_t SymbolLocation::Position::MaxLine;
 constexpr uint32_t SymbolLocation::Position::MaxColumn;
 
 void SymbolLocation::Position::setLine(uint32_t L) {
-  if (L > MaxLine) {
-    Line = MaxLine;
-    return;
-  }
-  Line = L;
+  if (L > MaxLine)
+    L = MaxLine;
+  LineColumnPacked = (L << ColumnBits) | column();
 }
 void SymbolLocation::Position::setColumn(uint32_t Col) {
-  if (Col > MaxColumn) {
-    Column = MaxColumn;
-    return;
-  }
-  Column = Col;
+  if (Col > MaxColumn)
+    Col = MaxColumn;
+  LineColumnPacked = (LineColumnPacked & ~MaxColumn) | Col;
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const SymbolLocation &L) {
   if (!L)
     return OS << "(none)";
   return OS << L.FileURI << "[" << L.Start.line() << ":" << L.Start.column()
             << "-" << L.End.line() << ":" << L.End.column() << ")";
 }
 
 } // namespace clangd
 } // namespace clang