Fix toHalfOpenFileRange assertion fail

[SureYeaah]

Summary

• Removed assumption that begin and end need to be in same file id.
• Fixes the crash at https://github.com/clangd/clangd/issues/113

[SureYeaah]

Implemented function to get expansion range such that ends are both in the same file id.

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

307

nit: getExpansionTokenRangeInSameFile?

315

Having a little trouble parsing this comment...

"Record the stack of expansion locations for the beginning, keyed by FileID"?

316

DenseMap<FileID, SourceRange>

Why are the values SourceRange here and not just SourceLocation?

318

Is there a way to write this loop that better captures the intent? while(1) with a break in the middle is hard for readers to understand at a glance.

I think it's just

for (SourceLocation Begin = ExpansionRange.getBegin();
       Begin.isValid();
       Begin = Begin.isFileId() ? SourceLocation() : SM.getImmediateExpansionRange(Begin).getBegin())
  BeginExpansions[SM.getFileID(Begin)] = Begin;

323

this call to toTokenRange seems unneccesary, that just adjusts End, but we only care about Start

328

again, I think by treating these as ranges rather than points you're making it harder than it needs to be. Traversing up the expansion tree isn't going to invert the order (though spelling might)

333

isn't this always SourceRange(It->second.getBegin(), EndExpansion.getEnd())?

357

nit: as with comments, avoid adding messages to the assertion unless it adds something to the code

[SureYeaah]

Addressed review comments

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

316

meant to say here: I think FileID is a legitimate key in a DenseMap, no need for unsigned

(the DenseMapInfo<clang::FileID> specialization exists, which should make it work)

[SureYeaah]

use fileID instead of hash in map

[sammccall]

Sorry, should have accepted this in the last round. Great stuff, thank you!

clang-tools-extra/clangd/SourceCode.cpp

static SourceRange unionTokenRange(SourceRange R1, SourceRange R2,
                                    const SourceManager &SM,
                                    const LangOptions &LangOpts) {
   SourceLocation E1 = getLocForTokenEnd(R1.getEnd(), SM, LangOpts);
   SourceLocation E2 = getLocForTokenEnd(R2.getEnd(), SM, LangOpts);
   return SourceRange(std::min(R1.getBegin(), R2.getBegin()),
                      E1 < E2 ? R2.getEnd() : R1.getEnd());
 }
 
-// Returns the tokenFileRange for a given Location as a Token Range
+// Check if two locations have the same file id.
+static bool inSameFile(SourceLocation Loc1, SourceLocation Loc2,
+                       const SourceManager &SM) {
+  return SM.getFileID(Loc1) == SM.getFileID(Loc2);
+}
+
+// Find an expansion range (not necessarily immediate) the ends of which are in
+// the same file id.
+static SourceRange
+getExpansionTokenRangeInSameFile(SourceLocation Loc, const SourceManager &SM,
+                                 const LangOptions &LangOpts) {
+  SourceRange ExpansionRange =
+      toTokenRange(SM.getImmediateExpansionRange(Loc), SM, LangOpts);
+  // Fast path for most common cases.
+  if (inSameFile(ExpansionRange.getBegin(), ExpansionRange.getEnd(), SM))
+    return ExpansionRange;
+  // Record the stack of expansion locations for the beginning, keyed by FileID.
+  llvm::DenseMap<unsigned, SourceLocation> BeginExpansions;
+  for (SourceLocation Begin = ExpansionRange.getBegin(); Begin.isValid();
+       Begin = Begin.isFileID()
+                   ? SourceLocation()
+                   : SM.getImmediateExpansionRange(Begin).getBegin()) {
+    BeginExpansions[SM.getFileID(Begin).getHashValue()] = Begin;
+  }
+  // Move up the stack of expansion locations for the end until we find the
+  // location in BeginExpansions with that has the same file id.
+  for (SourceLocation End = ExpansionRange.getEnd(); End.isValid();
+       End = End.isFileID() ? SourceLocation()
+                            : toTokenRange(SM.getImmediateExpansionRange(End),
+                                           SM, LangOpts)
+                                  .getEnd()) {
+    auto It = BeginExpansions.find(SM.getFileID(End).getHashValue());
+    if (It != BeginExpansions.end())
+      return {It->second, End};
+  }
+  llvm_unreachable(
+      "We should able to find a common ancestor in the expansion tree.");
+}
+// Returns the file range for a given Location as a Token Range
 // This is quite similar to getFileLoc in SourceManager as both use
 // getImmediateExpansionRange and getImmediateSpellingLoc (for macro IDs).
 // However:
 // - We want to maintain the full range information as we move from one file to
 //   the next. getFileLoc only uses the BeginLoc of getImmediateExpansionRange.
-// - We want to split '>>' tokens as the lexer parses the '>>' in template
-//   instantiations as a '>>' instead of a '>'.
+// - We want to split '>>' tokens as the lexer parses the '>>' in nested
+//   template instantiations as a '>>' instead of two '>'s.
 // There is also getExpansionRange but it simply calls
 // getImmediateExpansionRange on the begin and ends separately which is wrong.
 static SourceRange getTokenFileRange(SourceLocation Loc,
                                      const SourceManager &SM,
                                      const LangOptions &LangOpts) {
   SourceRange FileRange = Loc;
   while (!FileRange.getBegin().isFileID()) {
-    assert(!FileRange.getEnd().isFileID() &&
-           "Both Begin and End should be MacroIDs.");
     if (SM.isMacroArgExpansion(FileRange.getBegin())) {
-      FileRange.setBegin(SM.getImmediateSpellingLoc(FileRange.getBegin()));
-      FileRange.setEnd(SM.getImmediateSpellingLoc(FileRange.getEnd()));
+      FileRange = unionTokenRange(
+          SM.getImmediateSpellingLoc(FileRange.getBegin()),
+          SM.getImmediateSpellingLoc(FileRange.getEnd()), SM, LangOpts);
+      assert(inSameFile(FileRange.getBegin(), FileRange.getEnd(), SM));
     } else {
-      SourceRange ExpansionRangeForBegin = toTokenRange(
-          SM.getImmediateExpansionRange(FileRange.getBegin()), SM, LangOpts);
-      SourceRange ExpansionRangeForEnd = toTokenRange(
-          SM.getImmediateExpansionRange(FileRange.getEnd()), SM, LangOpts);
+      SourceRange ExpansionRangeForBegin =
+          getExpansionTokenRangeInSameFile(FileRange.getBegin(), SM, LangOpts);
+      SourceRange ExpansionRangeForEnd =
+          getExpansionTokenRangeInSameFile(FileRange.getEnd(), SM, LangOpts);
+      assert(inSameFile(ExpansionRangeForBegin.getBegin(),
+                        ExpansionRangeForEnd.getBegin(), SM) &&
+             "Both Expansion ranges should be in same file.");
       FileRange = unionTokenRange(ExpansionRangeForBegin, ExpansionRangeForEnd,
                                   SM, LangOpts);
     }
   }
   return FileRange;
 }
 
 bool isInsideMainFile(SourceLocation Loc, const SourceManager &SM) {

clang-tools-extra/clangd/unittests/SourceCodeTests.cpp

 // Test for functions toHalfOpenFileRange and getHalfOpenFileRange
 TEST(SourceCodeTests, HalfOpenFileRange) {
   // Each marked range should be the file range of the decl with the same name
   // and each name should be unique.
   Annotations Test(R"cpp(
     #define FOO(X, Y) int Y = ++X
     #define BAR(X) X + 1
     #define ECHO(X) X
+
+    #define BUZZ BAZZ(ADD)
+    #define BAZZ(m) m(1)
+    #define ADD(a) int f = a + 1;
     template<typename T>
     class P {};
-    void f() {
+
+    int main() {
       $a[[P<P<P<P<P<int>>>>> a]];
       $b[[int b = 1]];
       $c[[FOO(b, c)]]; 
       $d[[FOO(BAR(BAR(b)), d)]];
       // FIXME: We might want to select everything inside the outer ECHO.
       ECHO(ECHO($e[[int) ECHO(e]]));
+      // Shouldn't crash.
+      $f[[BUZZ]];
     }
   )cpp");
 
   ParsedAST AST = TestTU::withCode(Test.code()).build();
   llvm::errs() << Test.code();
   const SourceManager &SM = AST.getSourceManager();
   const LangOptions &LangOpts = AST.getASTContext().getLangOpts();
   // Turn a SourceLocation into a pair of positions
auto CheckRange = [&](llvm::StringRef Name) {
     EXPECT_EQ(HalfOpenRange, Test.ranges(Name)[0]);
   };
 
   CheckRange("a");
   CheckRange("b");
   CheckRange("c");
   CheckRange("d");
   CheckRange("e");
+  CheckRange("f");
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang