[clangd] Simplify the code in Index::refs

Summary:
While here, also fix potential UB in MergeIndex.

Thanks Kadir for finding this!

Reviewers: hokein

Reviewed By: hokein

Subscribers: merge_guards_bot, MaskRay, jkorous, arphaman, kadircet, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D70225

3 files changed, 17 insertions, 26 deletions

diff --git a/clang-tools-extra/clangd/index/MemIndex.cpp b/clang-tools-extra/clangd/index/MemIndex.cpp
index ea691b46cf3..71abca91cf0 100644
--- a/clang-tools-extra/clangd/index/MemIndex.cpp
+++ b/clang-tools-extra/clangd/index/MemIndex.cpp
@@ -72,7 +72,6 @@ bool MemIndex::refs(const RefsRequest &Req,
   trace::Span Tracer("MemIndex refs");
   uint32_t Remaining =
       Req.Limit.getValueOr(std::numeric_limits<uint32_t>::max());
-  bool More = false;
   for (const auto &ReqID : Req.IDs) {
     auto SymRefs = Refs.find(ReqID);
     if (SymRefs == Refs.end())
@@ -80,17 +79,13 @@ bool MemIndex::refs(const RefsRequest &Req,
     for (const auto &O : SymRefs->second) {
       if (!static_cast<int>(Req.Filter & O.Kind))
         continue;
-      if (Remaining == 0) {
-        More = true;
-        break;
-      }
-      if (Remaining > 0) {
-        --Remaining;
-        Callback(O);
-      }
+      if (Remaining == 0)
+        return true; // More refs were available.
+      --Remaining;
+      Callback(O);
     }
   }
-  return More;
+  return false; // We reported all refs.
 }
 
 void MemIndex::relations(
diff --git a/clang-tools-extra/clangd/index/Merge.cpp b/clang-tools-extra/clangd/index/Merge.cpp
index 8d8150a1313..0cef7dc7630 100644
--- a/clang-tools-extra/clangd/index/Merge.cpp
+++ b/clang-tools-extra/clangd/index/Merge.cpp
@@ -107,23 +107,24 @@ bool MergedIndex::refs(const RefsRequest &Req,
   More |= Dynamic->refs(Req, [&](const Ref &O) {
     DynamicIndexFileURIs.insert(O.Location.FileURI);
     Callback(O);
+    assert(Remaining != 0);
     --Remaining;
   });
   if (Remaining == 0 && More)
     return More;
   // We return less than Req.Limit if static index returns more refs for dirty
   // files.
-  More |= Static->refs(Req, [&](const Ref &O) {
+  bool StaticHadMore =  Static->refs(Req, [&](const Ref &O) {
     if (DynamicIndexFileURIs.count(O.Location.FileURI))
       return; // ignore refs that have been seen from dynamic index.
-    if (Remaining == 0)
+    if (Remaining == 0) {
       More = true;
-    if (Remaining > 0) {
-      --Remaining;
-      Callback(O);
+      return;
     }
+    --Remaining;
+    Callback(O);
   });
-  return More;
+  return More || StaticHadMore;
 }
 
 void MergedIndex::relations(
diff --git a/clang-tools-extra/clangd/index/dex/Dex.cpp b/clang-tools-extra/clangd/index/dex/Dex.cpp
index 3b4ed71da20..36ddedd04e1 100644
--- a/clang-tools-extra/clangd/index/dex/Dex.cpp
+++ b/clang-tools-extra/clangd/index/dex/Dex.cpp
@@ -254,21 +254,16 @@ bool Dex::refs(const RefsRequest &Req,
   trace::Span Tracer("Dex refs");
   uint32_t Remaining =
       Req.Limit.getValueOr(std::numeric_limits<uint32_t>::max());
-  bool More = false;
   for (const auto &ID : Req.IDs)
     for (const auto &Ref : Refs.lookup(ID)) {
       if (!static_cast<int>(Req.Filter & Ref.Kind))
         continue;
-      if (Remaining == 0) {
-        More = true;
-        break;
-      }
-      if (Remaining > 0) {
-        --Remaining;
-        Callback(Ref);
-      }
+      if (Remaining == 0)
+        return true; // More refs were available.
+      --Remaining;
+      Callback(Ref);
     }
-  return More;
+  return false; // We reported all refs.
 }
 
 void Dex::relations(