Enhance Lexer::makeFileCharRange to check for ranges inside a macro argument

expansion, in which case it returns a file range in the location where the
argument was spelled.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@148551 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 95 insertions, 13 deletions

diff --git a/lib/Lex/Lexer.cpp b/lib/Lex/Lexer.cpp
index 12cb76722b..967359f261 100644
--- a/lib/Lex/Lexer.cpp
+++ b/lib/Lex/Lexer.cpp
@@ -30,6 +30,7 @@
 #include "clang/Lex/CodeCompletionHandler.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <cstring>
@@ -792,6 +793,30 @@ bool Lexer::isAtEndOfMacroExpansion(SourceLocation loc,
   return isAtEndOfMacroExpansion(expansionLoc, SM, LangOpts, MacroEnd);
 }
 
+static CharSourceRange makeRangeFromFileLocs(SourceLocation Begin,
+                                             SourceLocation End,
+                                             const SourceManager &SM,
+                                             const LangOptions &LangOpts) {
+  assert(Begin.isFileID() && End.isFileID());
+  End = Lexer::getLocForEndOfToken(End, 0, SM,LangOpts);
+  if (End.isInvalid())
+    return CharSourceRange();
+
+  // Break down the source locations.
+  FileID FID;
+  unsigned BeginOffs;
+  llvm::tie(FID, BeginOffs) = SM.getDecomposedLoc(Begin);
+  if (FID.isInvalid())
+    return CharSourceRange();
+
+  unsigned EndOffs;
+  if (!SM.isInFileID(End, FID, &EndOffs) ||
+      BeginOffs > EndOffs)
+    return CharSourceRange();
+
+  return CharSourceRange::getCharRange(Begin, End);
+}
+
 /// \brief Accepts a token source range and returns a character range with
 /// file locations.
 /// Returns a null range if a part of the range resides inside a macro
@@ -800,28 +825,53 @@ CharSourceRange Lexer::makeFileCharRange(SourceRange TokenRange,
                                          const SourceManager &SM,
                                          const LangOptions &LangOpts) {
   SourceLocation Begin = TokenRange.getBegin();
-  if (Begin.isInvalid())
+  SourceLocation End = TokenRange.getEnd();
+  if (Begin.isInvalid() || End.isInvalid())
     return CharSourceRange();
 
-  if (Begin.isMacroID())
+  if (Begin.isFileID() && End.isFileID())
+    return makeRangeFromFileLocs(Begin, End, SM, LangOpts);
+
+  if (Begin.isMacroID() && End.isFileID()) {
     if (!isAtStartOfMacroExpansion(Begin, SM, LangOpts, &Begin))
       return CharSourceRange();
+    return makeRangeFromFileLocs(Begin, End, SM, LangOpts);
+  }
 
-  SourceLocation End = getLocForEndOfToken(TokenRange.getEnd(), 0, SM,LangOpts);
-  if (End.isInvalid())
-    return CharSourceRange();
+  if (Begin.isFileID() && End.isMacroID()) {
+    if (!isAtEndOfMacroExpansion(End, SM, LangOpts, &End))
+      return CharSourceRange();
+    return makeRangeFromFileLocs(Begin, End, SM, LangOpts);
+  }
 
-  // Break down the source locations.
-  std::pair<FileID, unsigned> beginInfo = SM.getDecomposedLoc(Begin);
-  if (beginInfo.first.isInvalid())
+  assert(Begin.isMacroID() && End.isMacroID());
+  SourceLocation MacroBegin, MacroEnd;
+  if (isAtStartOfMacroExpansion(Begin, SM, LangOpts, &MacroBegin) &&
+      isAtEndOfMacroExpansion(End, SM, LangOpts, &MacroEnd))
+    return makeRangeFromFileLocs(MacroBegin, MacroEnd, SM, LangOpts);
+
+  FileID FID;
+  unsigned BeginOffs;
+  llvm::tie(FID, BeginOffs) = SM.getDecomposedLoc(Begin);
+  if (FID.isInvalid())
     return CharSourceRange();
 
   unsigned EndOffs;
-  if (!SM.isInFileID(End, beginInfo.first, &EndOffs) ||
-      beginInfo.second > EndOffs)
+  if (!SM.isInFileID(End, FID, &EndOffs) ||
+      BeginOffs > EndOffs)
     return CharSourceRange();
 
-  return CharSourceRange::getCharRange(Begin, End);
+  const SrcMgr::SLocEntry *E = &SM.getSLocEntry(FID);
+  const SrcMgr::ExpansionInfo &Expansion = E->getExpansion();
+  if (Expansion.isMacroArgExpansion() &&
+      Expansion.getSpellingLoc().isFileID()) {
+    SourceLocation SpellLoc = Expansion.getSpellingLoc();
+    return makeRangeFromFileLocs(SpellLoc.getLocWithOffset(BeginOffs),
+                                 SpellLoc.getLocWithOffset(EndOffs),
+                                 SM, LangOpts);
+  }
+
+  return CharSourceRange();
 }
 
 StringRef Lexer::getSourceText(CharSourceRange Range,
diff --git a/unittests/Lex/LexerTest.cpp b/unittests/Lex/LexerTest.cpp
index c83a6ac8cf..7d9e9e9cbc 100644
--- a/unittests/Lex/LexerTest.cpp
+++ b/unittests/Lex/LexerTest.cpp
@@ -58,7 +58,8 @@ class VoidModuleLoader : public ModuleLoader {
 TEST_F(LexerTest, LexAPI) {
   const char *source =
     "#define M(x) [x]\n"
-    "M(foo)";
+    "#define N(x) x\n"
+    "M(foo) N([bar])";
   MemoryBuffer *buf = MemoryBuffer::getMemBuffer(source);
   SourceMgr.createMainFileIDForMemBuffer(buf);
 
@@ -82,10 +83,13 @@ TEST_F(LexerTest, LexAPI) {
   }
 
   // Make sure we got the tokens that we expected.
-  ASSERT_EQ(3U, toks.size());
+  ASSERT_EQ(6U, toks.size());
   ASSERT_EQ(tok::l_square, toks[0].getKind());
   ASSERT_EQ(tok::identifier, toks[1].getKind());
   ASSERT_EQ(tok::r_square, toks[2].getKind());
+  ASSERT_EQ(tok::l_square, toks[3].getKind());
+  ASSERT_EQ(tok::identifier, toks[4].getKind());
+  ASSERT_EQ(tok::r_square, toks[5].getKind());
   
   SourceLocation lsqrLoc = toks[0].getLocation();
   SourceLocation idLoc = toks[1].getLocation();
@@ -119,6 +123,34 @@ TEST_F(LexerTest, LexAPI) {
                   CharSourceRange::getTokenRange(SourceRange(lsqrLoc, rsqrLoc)),
                   SourceMgr, LangOpts);
   EXPECT_EQ(text, "M(foo)");
+
+  SourceLocation macroLsqrLoc = toks[3].getLocation();
+  SourceLocation macroIdLoc = toks[4].getLocation();
+  SourceLocation macroRsqrLoc = toks[5].getLocation();
+  SourceLocation fileLsqrLoc = SourceMgr.getSpellingLoc(macroLsqrLoc);
+  SourceLocation fileIdLoc = SourceMgr.getSpellingLoc(macroIdLoc);
+  SourceLocation fileRsqrLoc = SourceMgr.getSpellingLoc(macroRsqrLoc);
+
+  range = Lexer::makeFileCharRange(SourceRange(macroLsqrLoc, macroIdLoc),
+                                   SourceMgr, LangOpts);
+  EXPECT_EQ(SourceRange(fileLsqrLoc, fileIdLoc.getLocWithOffset(3)),
+            range.getAsRange());
+
+  range = Lexer::makeFileCharRange(SourceRange(macroIdLoc, macroRsqrLoc),
+                                   SourceMgr, LangOpts);
+  EXPECT_EQ(SourceRange(fileIdLoc, fileRsqrLoc.getLocWithOffset(1)),
+            range.getAsRange());
+
+  macroPair = SourceMgr.getExpansionRange(macroLsqrLoc);
+  range = Lexer::makeFileCharRange(SourceRange(macroLsqrLoc, macroRsqrLoc),
+                                   SourceMgr, LangOpts);
+  EXPECT_EQ(SourceRange(macroPair.first, macroPair.second.getLocWithOffset(1)),
+            range.getAsRange());
+
+  text = Lexer::getSourceText(
+          CharSourceRange::getTokenRange(SourceRange(macroLsqrLoc, macroIdLoc)),
+          SourceMgr, LangOpts);
+  EXPECT_EQ(text, "[bar");
 }
 
 } // anonymous namespace