Fixed toHalfOpenFileRange assertion fail

[SureYeaah]

Summary

• Added new function that gets Expansion range with both ends in the

same file.

• Fixes the crash at https://github.com/clangd/clangd/issues/113

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