[clangd] Don't trim xrefs references if we overran the limit

[kadircet]

Summary

This preserves all the results we've processed already rather than
throwing them away in the end.
It has some performance implications on the edge cases, in the worst case we
might issue 1 relations and 2 xrefs requests in extra to deduce HasMore
correctly.

Fixes https://github.com/clangd/clangd/issues/204.

[sammccall]

clang-tools-extra/clangd/XRefs.cpp

1388

This is a nice optimization.
Mention it in the patch description?

[sammccall]

clang-tools-extra/clangd/XRefs.cpp

1388

Oops, it came from elsewhere!

clang-tools-extra/clangd/XRefs.cpp

if (auto ID = getSymbolID(Base))
       OverriddenMethods.insert(ID);
     getOverriddenMethods(Base, OverriddenMethods);
   }
 }
 } // namespace
 
 ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
                                 const SymbolIndex *Index) {
-  if (!Limit)
-    Limit = std::numeric_limits<uint32_t>::max();
   ReferencesResult Results;
   const SourceManager &SM = AST.getSourceManager();
   auto MainFilePath =
       getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
   if (!MainFilePath) {
     elog("Failed to get a path for the main file, so no references");
     return Results;
   }
   auto URIMainFile = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
 
-  llvm::DenseSet<SymbolID> IDs, OverriddenMethods;
+  llvm::DenseSet<SymbolID> IDsToQuery, OverriddenMethods;
 
   const auto *IdentifierAtCursor =
       syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens());
   llvm::Optional<DefinedMacro> Macro;
   if (IdentifierAtCursor)
     Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor());
   if (Macro) {
     // Handle references to macro.
     if (auto MacroSID = getSymbolID(Macro->Name, Macro->Info, SM)) {
       // Collect macro references from main file.
       const auto &IDToRefs = AST.getMacros().MacroRefs;
       auto Refs = IDToRefs.find(MacroSID);
       if (Refs != IDToRefs.end()) {
         for (const auto &Ref : Refs->second) {
           ReferencesResult::Reference Result;
           Result.Loc.range = Ref.Rng;
           Result.Loc.uri = URIMainFile;
           if (Ref.IsDefinition) {
             Result.Attributes |= ReferencesResult::Declaration;
             Result.Attributes |= ReferencesResult::Definition;
           }
           Results.References.push_back(std::move(Result));
         }
       }
-      IDs.insert(MacroSID);
+      IDsToQuery.insert(MacroSID);
     }
   } else {
     // Handle references to Decls.
 
     DeclRelationSet Relations =
         DeclRelation::TemplatePattern | DeclRelation::Alias;
     std::vector<const NamedDecl *> Decls =
         getDeclAtPosition(AST, *CurLoc, Relations);
-    llvm::DenseSet<SymbolID> Targets;
-    for (const NamedDecl *D : Decls)
-      if (auto ID = getSymbolID(D))
-        Targets.insert(ID);
+    llvm::DenseSet<SymbolID> TargetsInMainFile;
+    for (const NamedDecl *D : Decls) {
+      auto ID = getSymbolID(D);
+      if (!ID)
+        continue;
+      TargetsInMainFile.insert(ID);
+      // Not all symbols can be referenced from outside (e.g. function-locals).
+      // TODO: we could skip TU-scoped symbols here (e.g. static functions) if
+      // we know this file isn't a header. The details might be tricky.
+      if (D->getParentFunctionOrMethod())
+        continue;
+      IDsToQuery.insert(ID);
+    }
 
     RelationsRequest OverriddenBy;
     if (Index) {
       OverriddenBy.Predicate = RelationKind::OverriddenBy;
       for (const NamedDecl *ND : Decls) {
         // Special case: For virtual methods, report decl/def of overrides and
         // references to all overridden methods in complete type hierarchy.
         if (const auto *CMD = llvm::dyn_cast<CXXMethodDecl>(ND)) {
           if (CMD->isVirtual()) {
             if (auto ID = getSymbolID(CMD))
               OverriddenBy.Subjects.insert(ID);
             getOverriddenMethods(CMD, OverriddenMethods);
           }
         }
       }
     }
 
     // We traverse the AST to find references in the main file.
-    auto MainFileRefs = findRefs(Targets, AST, /*PerToken=*/false);
+    auto MainFileRefs = findRefs(TargetsInMainFile, AST, /*PerToken=*/false);
     // We may get multiple refs with the same location and different Roles, as
     // cross-reference is only interested in locations, we deduplicate them
     // by the location to avoid emitting duplicated locations.
     MainFileRefs.erase(std::unique(MainFileRefs.begin(), MainFileRefs.end(),
                                    [](const ReferenceFinder::Reference &L,
                                       const ReferenceFinder::Reference &R) {
                                      return L.SpelledTok.location() ==
                                             R.SpelledTok.location();
                                    }),
                        MainFileRefs.end());
     for (const auto &Ref : MainFileRefs) {
       ReferencesResult::Reference Result;
       Result.Loc.range = Ref.range(SM);
       Result.Loc.uri = URIMainFile;
       if (Ref.Role & static_cast<unsigned>(index::SymbolRole::Declaration))
         Result.Attributes |= ReferencesResult::Declaration;
       // clang-index doesn't report definitions as declarations, but they are.
       if (Ref.Role & static_cast<unsigned>(index::SymbolRole::Definition))
         Result.Attributes |=
             ReferencesResult::Definition | ReferencesResult::Declaration;
       Results.References.push_back(std::move(Result));
     }
     // Add decl/def of overridding methods.
-    if (Index && Results.References.size() <= Limit &&
-        !OverriddenBy.Subjects.empty())
-      Index->relations(
-          OverriddenBy, [&](const SymbolID &Subject, const Symbol &Object) {
-            const auto LSPLocDecl =
-                toLSPLocation(Object.CanonicalDeclaration, *MainFilePath);
-            const auto LSPLocDef =
-                toLSPLocation(Object.Definition, *MainFilePath);
-            if (LSPLocDecl && LSPLocDecl != LSPLocDef) {
-              ReferencesResult::Reference Result;
-              Result.Loc = std::move(*LSPLocDecl);
-              Result.Attributes =
-                  ReferencesResult::Declaration | ReferencesResult::Override;
-              Results.References.push_back(std::move(Result));
-            }
-            if (LSPLocDef) {
-              ReferencesResult::Reference Result;
-              Result.Loc = std::move(*LSPLocDef);
-              Result.Attributes = ReferencesResult::Declaration |
-                                  ReferencesResult::Definition |
-                                  ReferencesResult::Override;
-              Results.References.push_back(std::move(Result));
-            }
-          });
-
-    if (Index && Results.References.size() <= Limit) {
-      for (const Decl *D : Decls) {
-        // Not all symbols can be referenced from outside (e.g.
-        // function-locals).
-        // TODO: we could skip TU-scoped symbols here (e.g. static functions) if
-        // we know this file isn't a header. The details might be tricky.
-        if (D->getParentFunctionOrMethod())
-          continue;
-        if (auto ID = getSymbolID(D))
-          IDs.insert(ID);
-      }
+    if (Index && !OverriddenBy.Subjects.empty()) {
+      Index->relations(OverriddenBy, [&](const SymbolID &Subject,
+                                         const Symbol &Object) {
+        if (Limit && Results.References.size() >= Limit) {
+          Results.HasMore = true;
+          return;
+        }
+        const auto LSPLocDef = toLSPLocation(Object.Definition, *MainFilePath);
     }
   }
   // Now query the index for references from other files.
   auto QueryIndex = [&](llvm::DenseSet<SymbolID> IDs, bool AllowAttributes,
                         bool AllowMainFileSymbols) {
+    if (IDs.empty() || !Index || Results.HasMore)
+      return;
     RefsRequest Req;
     Req.IDs = std::move(IDs);
-    Req.Limit = Limit;
-    if (Req.IDs.empty() || !Index || Results.References.size() > Limit)
-      return;
+    if (Limit) {
+      if (Limit < Results.References.size()) {
+        // We've already filled our quota, still check the index to correctly
+        // return the `HasMore` info.
+        Req.Limit = 0;
+      } else {
+        // Query index only for the remaining size.
+        Req.Limit = Limit - Results.References.size();
+      }
+    }
     Results.HasMore |= Index->refs(Req, [&](const Ref &R) {
-      // No need to continue process if we reach the limit.
-      if (Results.References.size() > Limit)
-        return;
       auto LSPLoc = toLSPLocation(R.Location, *MainFilePath);
       // Avoid indexed results for the main file - the AST is authoritative.
       if (!LSPLoc ||
           (!AllowMainFileSymbols && LSPLoc->uri.file() == *MainFilePath))
         return;
       ReferencesResult::Reference Result;
       Result.Loc = std::move(*LSPLoc);
       if (AllowAttributes) {
         if ((R.Kind & RefKind::Declaration) == RefKind::Declaration)
           Result.Attributes |= ReferencesResult::Declaration;
         // FIXME: our index should definitely store def | decl separately!
         if ((R.Kind & RefKind::Definition) == RefKind::Definition)
           Result.Attributes |=
               ReferencesResult::Declaration | ReferencesResult::Definition;
       }
       Results.References.push_back(std::move(Result));
     });
   };
-  QueryIndex(std::move(IDs), /*AllowAttributes=*/true,
+  QueryIndex(std::move(IDsToQuery), /*AllowAttributes=*/true,
              /*AllowMainFileSymbols=*/false);
   // For a virtual method: Occurrences of BaseMethod should be treated as refs
   // and not as decl/def. Allow symbols from main file since AST does not report
   // these.
   QueryIndex(std::move(OverriddenMethods), /*AllowAttributes=*/false,
              /*AllowMainFileSymbols=*/true);
-  if (Results.References.size() > Limit) {
-    Results.HasMore = true;
-    Results.References.resize(Limit);
-  }
   return Results;
 }
 
 std::vector<SymbolDetails> getSymbolInfo(ParsedAST &AST, Position Pos) {
   const SourceManager &SM = AST.getSourceManager();
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());