[Lexer] Fix bug in `makeFileCharRange` called on split tokens.

[ymandel]

Summary

When the end loc of the specified range is a split token, makeFileCharRange
does not process it correctly. This patch adds proper support for split tokens.

[ymandel]

Gentle ping...

[sammccall]

Sorry about this. I stared at this a few times and it *seems* right to me, and fixes the bug...

clang/lib/Lex/Lexer.cpp

static CharSourceRange makeRangeFromFileLocs(CharSourceRange Range,
   unsigned EndOffs;
   if (!SM.isInFileID(End, FID, &EndOffs) ||
       BeginOffs > EndOffs)
     return {};
 
   return CharSourceRange::getCharRange(Begin, End);
 }
 
+// Assumes that `Loc` is in an expansion.
+static bool isInExpansionTokenRange(const SourceLocation Loc,
+                                    const SourceManager &SM) {
+  return SM.getSLocEntry(SM.getFileID(Loc))
+      .getExpansion()
+      .isExpansionTokenRange();
+}
+
 CharSourceRange Lexer::makeFileCharRange(CharSourceRange Range,
                                          const SourceManager &SM,
                                          const LangOptions &LangOpts) {
   SourceLocation Begin = Range.getBegin();
   SourceLocation End = Range.getEnd();
   if (Begin.isInvalid() || End.isInvalid())
     return {};
 
   if (Begin.isFileID() && End.isFileID())
     return makeRangeFromFileLocs(Range, SM, LangOpts);
 
   if (Begin.isMacroID() && End.isFileID()) {
     if (!isAtStartOfMacroExpansion(Begin, SM, LangOpts, &Begin))
       return {};
     Range.setBegin(Begin);
     return makeRangeFromFileLocs(Range, SM, LangOpts);
   }
 
   if (Begin.isFileID() && End.isMacroID()) {
-    if ((Range.isTokenRange() && !isAtEndOfMacroExpansion(End, SM, LangOpts,
-                                                          &End)) ||
-        (Range.isCharRange() && !isAtStartOfMacroExpansion(End, SM, LangOpts,
-                                                           &End)))
+    if (Range.isTokenRange()) {
+      if (!isAtEndOfMacroExpansion(End, SM, LangOpts, &End))
+        return {};
+      // Use the *original* end, not the expanded one in `End`.
+      Range.setTokenRange(isInExpansionTokenRange(Range.getEnd(), SM));
+    } else if (!isAtStartOfMacroExpansion(End, SM, LangOpts, &End))
       return {};
     Range.setEnd(End);
     return makeRangeFromFileLocs(Range, SM, LangOpts);
   }
 
   assert(Begin.isMacroID() && End.isMacroID());
   SourceLocation MacroBegin, MacroEnd;
   if (isAtStartOfMacroExpansion(Begin, SM, LangOpts, &MacroBegin) &&
       ((Range.isTokenRange() && isAtEndOfMacroExpansion(End, SM, LangOpts,
                                                         &MacroEnd)) ||
        (Range.isCharRange() && isAtStartOfMacroExpansion(End, SM, LangOpts,
                                                          &MacroEnd)))) {
     Range.setBegin(MacroBegin);
     Range.setEnd(MacroEnd);
+    // Use the *original* `End`, not the expanded one in `MacroEnd`.
+    if (Range.isTokenRange())
+      Range.setTokenRange(isInExpansionTokenRange(End, SM));
     return makeRangeFromFileLocs(Range, SM, LangOpts);
   }
 
   bool Invalid = false;
   const SrcMgr::SLocEntry &BeginEntry = SM.getSLocEntry(SM.getFileID(Begin),
                                                         &Invalid);
   if (Invalid)
     return {};

clang/unittests/Lex/LexerTest.cpp

 #include "clang/Lex/HeaderSearchOptions.h"
 #include "clang/Lex/MacroArgs.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/ModuleLoader.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <memory>
 #include <vector>
 
 namespace {
 using namespace clang;
 using testing::ElementsAre;
 
 // The test fixture.
 class LexerTest : public ::testing::Test {
std::unique_ptr<Preprocessor> PP = std::make_unique<Preprocessor>(
         /*OwnsHeaderSearch =*/false);
     PP->Initialize(*Target);
     PP->EnterMainSourceFile();
     return PP;
   }
 
   std::vector<Token> Lex(StringRef Source) {
     TrivialModuleLoader ModLoader;
-    auto PP = CreatePP(Source, ModLoader);
+    PP = CreatePP(Source, ModLoader);
 
     std::vector<Token> toks;
     while (1) {
       Token tok;
       PP->Lex(tok);
       if (tok.is(tok::eof))
         break;
       toks.push_back(tok);
protected:
   FileSystemOptions FileMgrOpts;
   FileManager FileMgr;
   IntrusiveRefCntPtr<DiagnosticIDs> DiagID;
   DiagnosticsEngine Diags;
   SourceManager SourceMgr;
   LangOptions LangOpts;
   std::shared_ptr<TargetOptions> TargetOpts;
   IntrusiveRefCntPtr<TargetInfo> Target;
+  std::unique_ptr<Preprocessor> PP;
 };
 
 TEST_F(LexerTest, GetSourceTextExpandsToMaximumInMacroArgument) {
   std::vector<tok::TokenKind> ExpectedTokens;
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::l_paren);
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::r_paren);
std::vector<Token> toks = CheckLex("#define M(x) x\n"
                                      "#define C(x) c M(x)\n"
                                      "C(f(M(i)))",
                                      ExpectedTokens);
   EXPECT_EQ("M(i)", getSourceText(toks[3], toks[3]));
 }
 
 TEST_F(LexerTest, LexAPI) {
   std::vector<tok::TokenKind> ExpectedTokens;
+  // Line 1 (after the #defines)
   ExpectedTokens.push_back(tok::l_square);
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::r_square);
   ExpectedTokens.push_back(tok::l_square);
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::r_square);
+  // Line 2
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::identifier);
   ExpectedTokens.push_back(tok::identifier);
 
   std::vector<Token> toks = CheckLex("#define M(x) [x]\n"
                                      "#define N(x) x\n"
                                      "#define INN(x) x\n"
TEST_F(LexerTest, LexAPI) {
   SourceLocation idLoc3 = toks[8].getLocation();
   SourceLocation idLoc4 = toks[9].getLocation();
   EXPECT_EQ("INN", Lexer::getImmediateMacroName(idLoc1, SourceMgr, LangOpts));
   EXPECT_EQ("INN", Lexer::getImmediateMacroName(idLoc2, SourceMgr, LangOpts));
   EXPECT_EQ("NOF2", Lexer::getImmediateMacroName(idLoc3, SourceMgr, LangOpts));
   EXPECT_EQ("N", Lexer::getImmediateMacroName(idLoc4, SourceMgr, LangOpts));
 }
 
+TEST_F(LexerTest, HandlesSplitTokens) {
+  std::vector<tok::TokenKind> ExpectedTokens;
+  // Line 1 (after the #defines)
+  ExpectedTokens.push_back(tok::identifier);
+  ExpectedTokens.push_back(tok::less);
+  ExpectedTokens.push_back(tok::identifier);
+  ExpectedTokens.push_back(tok::less);
+  ExpectedTokens.push_back(tok::greatergreater);
+  // Line 2
+  ExpectedTokens.push_back(tok::identifier);
+  ExpectedTokens.push_back(tok::less);
+  ExpectedTokens.push_back(tok::identifier);
+  ExpectedTokens.push_back(tok::less);
+  ExpectedTokens.push_back(tok::greatergreater);
+
+  std::vector<Token> toks = CheckLex("#define TY ty\n"
+                                     "#define RANGLE ty<ty<>>\n"
+                                     "TY<ty<>>\n"
+                                     "RANGLE",
+                                     ExpectedTokens);
+
+  SourceLocation outerTyLoc = toks[0].getLocation();
+  SourceLocation innerTyLoc = toks[2].getLocation();
+  SourceLocation gtgtLoc = toks[4].getLocation();
+  // Split the token to simulate the action of the parser and force creation of
+  // an `ExpansionTokenRange`.
+  SourceLocation rangleLoc = PP->SplitToken(gtgtLoc, 1);
+
+  // Verify that it only captures the first greater-then and not the second one.
+  CharSourceRange range = Lexer::makeFileCharRange(
+      CharSourceRange::getTokenRange(innerTyLoc, rangleLoc), SourceMgr,
+      LangOpts);
+  EXPECT_TRUE(range.isCharRange());
+  EXPECT_EQ(range.getAsRange(),
+            SourceRange(innerTyLoc, gtgtLoc.getLocWithOffset(1)));
+
+  // Verify case where range begins in a macro expansion.
+  range = Lexer::makeFileCharRange(
+      CharSourceRange::getTokenRange(outerTyLoc, rangleLoc), SourceMgr,
+      LangOpts);
+  EXPECT_TRUE(range.isCharRange());
+  EXPECT_EQ(range.getAsRange(),
+            SourceRange(SourceMgr.getExpansionLoc(outerTyLoc),
+                        gtgtLoc.getLocWithOffset(1)));
+
+  SourceLocation macroInnerTyLoc = toks[7].getLocation();
+  SourceLocation macroGtgtLoc = toks[9].getLocation();
+  // Split the token to simulate the action of the parser and force creation of
+  // an `ExpansionTokenRange`.
+  SourceLocation macroRAngleLoc = PP->SplitToken(macroGtgtLoc, 1);
+
+  // Verify that it fails (because it only captures the first greater-then and
+  // not the second one, so it doesn't span the entire macro expansion).
+  range = Lexer::makeFileCharRange(
+      CharSourceRange::getTokenRange(macroInnerTyLoc, macroRAngleLoc),
+      SourceMgr, LangOpts);
+  EXPECT_TRUE(range.isInvalid());
+}
+
 TEST_F(LexerTest, DontMergeMacroArgsFromDifferentMacroFiles) {
   std::vector<Token> toks =
       Lex("#define helper1 0\n"
           "void helper2(const char *, ...);\n"
           "#define M1(a, ...) helper2(a, ##__VA_ARGS__)\n"
           "#define M2(a, ...) M1(a, helper1, ##__VA_ARGS__)\n"
           "void f1() { M2(\"a\", \"b\"); }");