1 files changed, 58 insertions, 42 deletions
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 05e04ac..ef45acf 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -21,6 +21,7 @@
 #include "clang-include-cleaner/Types.h"
 #include "index/Index.h"
 #include "index/Merge.h"
+#include "index/Ref.h"
 #include "index/Relation.h"
 #include "index/SymbolCollector.h"
 #include "index/SymbolID.h"
@@ -56,6 +57,7 @@
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallSet.h"
@@ -66,6 +68,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
+#include <algorithm>
 #include <optional>
 #include <string>
 #include <vector>
@@ -2350,51 +2353,64 @@ incomingCalls(const CallHierarchyItem &Item, const SymbolIndex *Index) {
   // to an AST node isn't cheap, particularly when the declaration isn't
   // in the main file.
   // FIXME: Consider also using AST information when feasible.
-  RefsRequest Request;
-  Request.IDs.insert(*ID);
-  Request.WantContainer = true;
-  // We could restrict more specifically to calls by introducing a new RefKind,
-  // but non-call references (such as address-of-function) can still be
-  // interesting as they can indicate indirect calls.
-  Request.Filter = RefKind::Reference;
-  // Initially store the ranges in a map keyed by SymbolID of the caller.
-  // This allows us to group different calls with the same caller
-  // into the same CallHierarchyIncomingCall.
-  llvm::DenseMap<SymbolID, std::vector<Location>> CallsIn;
-  // We can populate the ranges based on a refs request only. As we do so, we
-  // also accumulate the container IDs into a lookup request.
-  LookupRequest ContainerLookup;
-  Index->refs(Request, [&](const Ref &R) {
-    auto Loc = indexToLSPLocation(R.Location, Item.uri.file());
-    if (!Loc) {
-      elog("incomingCalls failed to convert location: {0}", Loc.takeError());
-      return;
-    }
-    CallsIn[R.Container].push_back(*Loc);
+  auto QueryIndex = [&](llvm::DenseSet<SymbolID> IDs, bool MightNeverCall) {
+    RefsRequest Request;
+    Request.IDs = std::move(IDs);
+    Request.WantContainer = true;
+    // We could restrict more specifically to calls by introducing a new
+    // RefKind, but non-call references (such as address-of-function) can still
+    // be interesting as they can indicate indirect calls.
+    Request.Filter = RefKind::Reference;
+    // Initially store the ranges in a map keyed by SymbolID of the caller.
+    // This allows us to group different calls with the same caller
+    // into the same CallHierarchyIncomingCall.
+    llvm::DenseMap<SymbolID, std::vector<Location>> CallsIn;
+    // We can populate the ranges based on a refs request only. As we do so, we
+    // also accumulate the container IDs into a lookup request.
+    LookupRequest ContainerLookup;
+    Index->refs(Request, [&](const Ref &R) {
+      auto Loc = indexToLSPLocation(R.Location, Item.uri.file());
+      if (!Loc) {
+        elog("incomingCalls failed to convert location: {0}", Loc.takeError());
+        return;
+      }
+      CallsIn[R.Container].push_back(*Loc);
 
-    ContainerLookup.IDs.insert(R.Container);
-  });
-  // Perform the lookup request and combine its results with CallsIn to
-  // get complete CallHierarchyIncomingCall objects.
-  Index->lookup(ContainerLookup, [&](const Symbol &Caller) {
-    auto It = CallsIn.find(Caller.ID);
-    assert(It != CallsIn.end());
-    if (auto CHI = symbolToCallHierarchyItem(Caller, Item.uri.file())) {
-      std::vector<Range> FromRanges;
-      for (const Location &L : It->second) {
-        if (L.uri != CHI->uri) {
-          // Call location not in same file as caller.
-          // This can happen in some edge cases. There's not much we can do,
-          // since the protocol only allows returning ranges interpreted as
-          // being in the caller's file.
-          continue;
+      ContainerLookup.IDs.insert(R.Container);
+    });
+    // Perform the lookup request and combine its results with CallsIn to
+    // get complete CallHierarchyIncomingCall objects.
+    Index->lookup(ContainerLookup, [&](const Symbol &Caller) {
+      auto It = CallsIn.find(Caller.ID);
+      assert(It != CallsIn.end());
+      if (auto CHI = symbolToCallHierarchyItem(Caller, Item.uri.file())) {
+        std::vector<Range> FromRanges;
+        for (const Location &L : It->second) {
+          if (L.uri != CHI->uri) {
+            // Call location not in same file as caller.
+            // This can happen in some edge cases. There's not much we can do,
+            // since the protocol only allows returning ranges interpreted as
+            // being in the caller's file.
+            continue;
+          }
+          FromRanges.push_back(L.range);
         }
-        FromRanges.push_back(L.range);
+        Results.push_back(CallHierarchyIncomingCall{
+            std::move(*CHI), std::move(FromRanges), MightNeverCall});
       }
-      Results.push_back(
-          CallHierarchyIncomingCall{std::move(*CHI), std::move(FromRanges)});
-    }
-  });
+    });
+  };
+  QueryIndex({ID.get()}, false);
+  // In the case of being a virtual function we also want to return
+  // potential calls through the base function.
+  if (Item.kind == SymbolKind::Method) {
+    llvm::DenseSet<SymbolID> IDs;
+    RelationsRequest Req{{ID.get()}, RelationKind::OverriddenBy, std::nullopt};
+    Index->reverseRelations(Req, [&](const SymbolID &, const Symbol &Caller) {
+      IDs.insert(Caller.ID);
+    });
+    QueryIndex(std::move(IDs), true);
+  }
   // Sort results by name of container.
   llvm::sort(Results, [](const CallHierarchyIncomingCall &A,
                          const CallHierarchyIncomingCall &B) {