[clang-tools-extra] [clangd] fix crash in include cleaner (PR #99514)

[Yuxuan Shui via cfe-commits]

https://github.com/yshui updated https://github.com/llvm/llvm-project/pull/99514

>From 6e5d8d8e04678d80b3139c6d6129da88df57c945 Mon Sep 17 00:00:00 2001
From: Yuxuan Shui <yshuiv7 at gmail.com>
Date: Thu, 18 Jul 2024 16:07:45 +0100
Subject: [PATCH] [clangd] [IncludeCleaner] Use correct file ID

---
 clang-tools-extra/clangd/CollectMacros.cpp  | 10 +++++-----
 clang-tools-extra/clangd/CollectMacros.h    |  1 +
 clang-tools-extra/clangd/IncludeCleaner.cpp |  2 +-
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/clang-tools-extra/clangd/CollectMacros.cpp b/clang-tools-extra/clangd/CollectMacros.cpp
index c5ba8d903ba48..83f145d25a5a6 100644
--- a/clang-tools-extra/clangd/CollectMacros.cpp
+++ b/clang-tools-extra/clangd/CollectMacros.cpp
@@ -19,9 +19,8 @@ namespace clang {
 namespace clangd {
 
 Range MacroOccurrence::toRange(const SourceManager &SM) const {
-  auto MainFile = SM.getMainFileID();
   return halfOpenToRange(
-      SM, syntax::FileRange(MainFile, StartOffset, EndOffset).toCharRange(SM));
+      SM, syntax::FileRange(FID, StartOffset, EndOffset).toCharRange(SM));
 }
 
 void CollectMainFileMacros::add(const Token &MacroNameTok, const MacroInfo *MI,
@@ -34,12 +33,13 @@ void CollectMainFileMacros::add(const Token &MacroNameTok, const MacroInfo *MI,
 
   auto Name = MacroNameTok.getIdentifierInfo()->getName();
   Out.Names.insert(Name);
-  size_t Start = SM.getFileOffset(Loc);
+  auto [FID, Start] = SM.getDecomposedLoc(Loc);
   size_t End = SM.getFileOffset(MacroNameTok.getEndLoc());
   if (auto SID = getSymbolID(Name, MI, SM))
-    Out.MacroRefs[SID].push_back({Start, End, IsDefinition, InIfCondition});
+    Out.MacroRefs[SID].push_back(
+        {FID, Start, End, IsDefinition, InIfCondition});
   else
-    Out.UnknownMacros.push_back({Start, End, IsDefinition, InIfCondition});
+    Out.UnknownMacros.push_back({FID, Start, End, IsDefinition, InIfCondition});
 }
 
 void CollectMainFileMacros::FileChanged(SourceLocation Loc, FileChangeReason,
diff --git a/clang-tools-extra/clangd/CollectMacros.h b/clang-tools-extra/clangd/CollectMacros.h
index e3900c08e5df7..4de49b1be3408 100644
--- a/clang-tools-extra/clangd/CollectMacros.h
+++ b/clang-tools-extra/clangd/CollectMacros.h
@@ -24,6 +24,7 @@ namespace clangd {
 
 struct MacroOccurrence {
   // Half-open range (end offset is exclusive) inside the main file.
+  FileID FID;
   size_t StartOffset;
   size_t EndOffset;
 
diff --git a/clang-tools-extra/clangd/IncludeCleaner.cpp b/clang-tools-extra/clangd/IncludeCleaner.cpp
index dc5b7ec95db5f..a739286abbf3d 100644
--- a/clang-tools-extra/clangd/IncludeCleaner.cpp
+++ b/clang-tools-extra/clangd/IncludeCleaner.cpp
@@ -302,7 +302,7 @@ collectMacroReferences(ParsedAST &AST) {
   std::vector<include_cleaner::SymbolReference> Macros;
   for (const auto &[_, Refs] : AST.getMacros().MacroRefs) {
     for (const auto &Ref : Refs) {
-      auto Loc = SM.getComposedLoc(SM.getMainFileID(), Ref.StartOffset);
+      auto Loc = SM.getComposedLoc(Ref.FID, Ref.StartOffset);
       const auto *Tok = AST.getTokens().spelledTokenContaining(Loc);
       if (!Tok)
         continue;