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

306

nit: getExpansionTokenRangeInSameFile?

314

Having a little trouble parsing this comment...

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

315

DenseMap<FileID, SourceRange>

Why are the values SourceRange here and not just SourceLocation?

317

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;

322

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

324

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

327

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)

332

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

[SureYeaah]

Addressed review comments

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

315

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/Selection.cpp

if (!S.isValid())
       return SelectionTree::Unselected;
     // toHalfOpenFileRange() allows selection of constructs in macro args. e.g:
     //   #define LOOP_FOREVER(Body) for(;;) { Body }
     //   void IncrementLots(int &x) {
     //     LOOP_FOREVER( ++x; )
     //   }
     // Selecting "++x" or "x" will do the right thing.
     auto Range = toHalfOpenFileRange(SM, LangOpts, S);
-    assert(Range && "We should be able to get the File Range");
+    if(!Range)
+      return SelectionTree::Unselected;
     dlog("{1}claimRange: {0}", Range->printToString(SM), indent());
     auto B = SM.getDecomposedLoc(Range->getBegin());
     auto E = SM.getDecomposedLoc(Range->getEnd());
     // Otherwise, nodes in macro expansions can't be selected.
     if (B.first != SelFile || E.first != SelFile)
       return SelectionTree::Unselected;
     // Attempt to claim the remaining range. If there's nothing to claim, only
     // children were selected.

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());
 }
+// If Loc is a macro arg expansion, return the immediate spelling loc.
+// Otherwise 
+// Assumes Loc is in a macro file ID
+static SourceRange getExpansionOrSpellingRange(SourceLocation Loc,
+                                               const SourceManager &SM,
+                                               const LangOptions &LangOpts) {
+  if (SM.isMacroArgExpansion(Loc))
+    return SM.getImmediateSpellingLoc(Loc);
+  return toTokenRange(SM.getImmediateExpansionRange(Loc), SM, LangOpts);
+}
 
-// Returns the tokenFileRange for a given Location as a Token Range
+// Returns the tokenFileRange for a given Range 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,
+static SourceRange getTokenFileRange(SourceRange Rng, 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()));
-    } else {
-      SourceRange ExpansionRangeForBegin = toTokenRange(
-          SM.getImmediateExpansionRange(FileRange.getBegin()), SM, LangOpts);
-      SourceRange ExpansionRangeForEnd = toTokenRange(
-          SM.getImmediateExpansionRange(FileRange.getEnd()), SM, LangOpts);
-      FileRange = unionTokenRange(ExpansionRangeForBegin, ExpansionRangeForEnd,
-                                  SM, LangOpts);
-    }
-  }
-  return FileRange;
+  // Map from FileID hash to SourceRange
+  std::map<unsigned, SourceRange> RngMap;
+  // Adds a Location to the RngMap
+  auto AddLocToMap = [&RngMap, &SM, &LangOpts](SourceLocation Loc) {
+    unsigned FileIDHash = SM.getFileID(Loc).getHashValue();
+    auto It = RngMap.find(FileIDHash);
+    if (It == RngMap.end())
+      RngMap[FileIDHash] = Loc;
+    else
+      It->second = unionTokenRange(It->second, Loc, SM, LangOpts);
+  };
+  auto AddRngToMap = [&](SourceRange SR) {
+    AddLocToMap(SR.getBegin());
+    AddLocToMap(SR.getEnd());
+  };
+  // Initialize RngMap
+  AddRngToMap(Rng);
+  // Iterate until only one range is left which is not a macroID
+  for (auto It = RngMap.rbegin();
+       !RngMap.empty() && It->second.getBegin().isMacroID();
+       It = RngMap.rbegin()) {
+    // Get SourceRange with maximum FileID and remove from RngMap
+    Rng = It->second;
+    RngMap.erase(It->first);
+    // Add Immediate Expansion/Spelling Range to RngMap
+    AddRngToMap(getExpansionOrSpellingRange(Rng.getBegin(), SM, LangOpts));
+    AddRngToMap(getExpansionOrSpellingRange(Rng.getEnd(), SM, LangOpts));
+  }
+  assert(RngMap.size() == 1 && "There should be only one final range.");
+  SourceRange TokenFileRange = RngMap.begin()->second;
+  assert(isInsideMainFile(TokenFileRange.getBegin(), SM) &&
+         isInsideMainFile(TokenFileRange.getBegin(), SM) &&
+         "Final range should be in main file.");
+  return TokenFileRange;
 }
 
 bool isInsideMainFile(SourceLocation Loc, const SourceManager &SM) {
   return Loc.isValid() && SM.isWrittenInMainFile(SM.getExpansionLoc(Loc));
 }
 
 llvm::Optional<SourceRange> toHalfOpenFileRange(const SourceManager &SM,
                                                 const LangOptions &LangOpts,
                                                 SourceRange R) {
-  SourceRange R1 = getTokenFileRange(R.getBegin(), SM, LangOpts);
-  if (!isValidFileRange(SM, R1))
-    return llvm::None;
-
-  SourceRange R2 = getTokenFileRange(R.getEnd(), SM, LangOpts);
-  if (!isValidFileRange(SM, R2))
+  SourceRange Result = getTokenFileRange(R, SM, LangOpts);
+  if (!isValidFileRange(SM, Result))
     return llvm::None;
 
-  SourceRange Result = unionTokenRange(R1, R2, SM, LangOpts);
   unsigned TokLen = getTokenLengthAtLoc(Result.getEnd(), SM, LangOpts);
   // Convert from closed token range to half-open (char) range
   Result.setEnd(Result.getEnd().getLocWithOffset(TokLen));
   if (!isValidFileRange(SM, Result))
     return llvm::None;
 
   return Result;
 }

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