[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 d8f59628643bc6eb4d572949d7f9df6df42f0146 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            | 21 +++++++++++---
 .../clangd/unittests/CallHierarchyTests.cpp   | 29 +++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 4fd11307857ff8..61fa66180376cd 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>
@@ -2275,7 +2276,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;
@@ -2285,7 +2286,7 @@ incomingCalls(const CallHierarchyItem &Item, const SymbolIndex *Index) {
       elog("incomingCalls failed to convert location: {0}", Loc.takeError());
       return;
     }
-    CallsIn[R.Container].push_back(Loc->range);
+    CallsIn[R.Container].push_back(*Loc);
 
     ContainerLookup.IDs.insert(R.Container);
   });
@@ -2294,9 +2295,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 b2278ff12735dc..8821d3aad9c784 100644
--- a/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp
+++ b/clang-tools-extra/clangd/unittests/CallHierarchyTests.cpp
@@ -491,6 +491,35 @@ TEST(CallHierarchy, HierarchyOnVar) {
                                 fromRanges(Source.range("Callee")))));
 }
 
+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