[clang-tools-extra] [clangd] Harden incomingCalls() against possible misinterpretation of a range as pertaining to the wrong file (PR #111616)

[Nathan Ridge via cfe-commits]

https://github.com/HighCommander4 updated https://github.com/llvm/llvm-project/pull/111616

>From c72cc8dedc165091654eb5d922927319327f34c4 Mon Sep 17 00:00:00 2001
From: Nathan Ridge <zeratul976 at hotmail.com>
Date: Tue, 8 Oct 2024 21:43:55 -0400
Subject: [PATCH] [clangd] Harden incomingCalls() against possible
 misinterpretation of a range as pertaining to the wrong file

---
 clang-tools-extra/clangd/XRefs.cpp            | 23 +++++++++++----
 .../clangd/unittests/CallHierarchyTests.cpp   | 29 +++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index f94cadeffaa298..424b6800f1960f 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -63,6 +63,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
@@ -2272,7 +2273,7 @@ incomingCalls(const CallHierarchyItem &Item, const SymbolIndex *Index) {
   // 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<Range>> CallsIn;
+  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;
@@ -2282,8 +2283,8 @@ incomingCalls(const CallHierarchyItem &Item, const SymbolIndex *Index) {
       elog("incomingCalls failed to convert location: {0}", Loc.takeError());
       return;
     }
-    auto It = CallsIn.try_emplace(R.Container, std::vector<Range>{}).first;
-    It->second.push_back(Loc->range);
+    auto It = CallsIn.try_emplace(R.Container, std::vector<Location>{}).first;
+    It->second.push_back(*Loc);
 
     ContainerLookup.IDs.insert(R.Container);
   });
@@ -2292,9 +2293,21 @@ incomingCalls(const CallHierarchyItem &Item, const SymbolIndex *Index) {
   Index->lookup(ContainerLookup, [&](const Symbol &Caller) {
     auto It = CallsIn.find(Caller.ID);
     assert(It != CallsIn.end());
-    if (auto CHI = symbolToCallHierarchyItem(Caller, Item.uri.file()))
+    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);
+      }
       Results.push_back(
-          CallHierarchyIncomingCall{std::move(*CHI), std::move(It->second)});
+          CallHierarchyIncomingCall{std::move(*CHI), std::move(FromRanges)});
+    }
   });
   // Sort results by name of container.
   llvm::sort(Results, [](const CallHierarchyIncomingCall &A,
diff --git a/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp b/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp
index 6fa76aa6094bf2..9518bc5041bcca 100644
--- a/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp
+++ b/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp
@@ -446,6 +446,35 @@ TEST(CallHierarchy, CallInLocalVarDecl) {
           AllOf(from(withName("caller3")), fromRanges(Source.range("call3")))));
 }
 
+TEST(CallHierarchy, CallInDifferentFileThanCaller) {
+  Annotations Header(R"cpp(
+    #define WALDO void caller() {
+  )cpp");
+  Annotations Source(R"cpp(
+    void call^ee();
+    WALDO
+      callee();
+    }
+  )cpp");
+  auto TU = TestTU::withCode(Source.code());
+  TU.HeaderCode = Header.code();
+  auto AST = TU.build();
+  auto Index = TU.index();
+
+  std::vector<CallHierarchyItem> Items =
+      prepareCallHierarchy(AST, Source.point(), testPath(TU.Filename));
+  ASSERT_THAT(Items, ElementsAre(withName("callee")));
+
+  auto Incoming = incomingCalls(Items[0], Index.get());
+
+  // The only call site is in the source file, which is a different file from
+  // the declaration of the function containing the call, which is in the
+  // header. The protocol does not allow us to represent such calls, so we drop
+  // them. (The call hierarchy item itself is kept.)
+  EXPECT_THAT(Incoming,
+              ElementsAre(AllOf(from(withName("caller")), fromRanges())));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang