[Lexer] Improve Lexer::getSourceText() when the given range deals with function macro arguments.

This is a modified version of a patch by Manuel Klimek.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@182055 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/lib/Basic/SourceManager.cpp b/lib/Basic/SourceManager.cpp
index d6dc6d6..6994b30 100644
--- a/lib/Basic/SourceManager.cpp
+++ b/lib/Basic/SourceManager.cpp
@@ -536,6 +536,43 @@
   return FakeContentCacheForRecovery;
 }
 
+/// \brief Returns the previous in-order FileID or an invalid FileID if there
+/// is no previous one.
+FileID SourceManager::getPreviousFileID(FileID FID) const {
+  if (FID.isInvalid())
+    return FileID();
+
+  int ID = FID.ID;
+  if (ID == -1)
+    return FileID();
+
+  if (ID > 0) {
+    if (ID-1 == 0)
+      return FileID();
+  } else if (unsigned(-(ID-1) - 2) >= LoadedSLocEntryTable.size()) {
+    return FileID();
+  }
+
+  return FileID::get(ID-1);
+}
+
+/// \brief Returns the next in-order FileID or an invalid FileID if there is
+/// no next one.
+FileID SourceManager::getNextFileID(FileID FID) const {
+  if (FID.isInvalid())
+    return FileID();
+
+  int ID = FID.ID;
+  if (ID > 0) {
+    if (unsigned(ID+1) >= local_sloc_entry_size())
+      return FileID();
+  } else if (ID+1 >= -1) {
+    return FileID();
+  }
+
+  return FileID::get(ID+1);
+}
+
 //===----------------------------------------------------------------------===//
 // Methods to create new FileID's and macro expansions.
 //===----------------------------------------------------------------------===//
@@ -998,6 +1035,77 @@
   return Expansion.isMacroBodyExpansion();
 }
 
+bool SourceManager::isAtStartOfImmediateMacroExpansion(SourceLocation Loc,
+                                             SourceLocation *MacroBegin) const {
+  assert(Loc.isValid() && Loc.isMacroID() && "Expected a valid macro loc");
+
+  std::pair<FileID, unsigned> DecompLoc = getDecomposedLoc(Loc);
+  if (DecompLoc.second > 0)
+    return false; // Does not point at the start of expansion range.
+
+  bool Invalid = false;
+  const SrcMgr::ExpansionInfo &ExpInfo =
+      getSLocEntry(DecompLoc.first, &Invalid).getExpansion();
+  if (Invalid)
+    return false;
+  SourceLocation ExpLoc = ExpInfo.getExpansionLocStart();
+
+  if (ExpInfo.isMacroArgExpansion()) {
+    // For macro argument expansions, check if the previous FileID is part of
+    // the same argument expansion, in which case this Loc is not at the
+    // beginning of the expansion.
+    FileID PrevFID = getPreviousFileID(DecompLoc.first);
+    if (!PrevFID.isInvalid()) {
+      const SrcMgr::SLocEntry &PrevEntry = getSLocEntry(PrevFID, &Invalid);
+      if (Invalid)
+        return false;
+      if (PrevEntry.isExpansion() &&
+          PrevEntry.getExpansion().getExpansionLocStart() == ExpLoc)
+        return false;
+    }
+  }
+
+  if (MacroBegin)
+    *MacroBegin = ExpLoc;
+  return true;
+}
+
+bool SourceManager::isAtEndOfImmediateMacroExpansion(SourceLocation Loc,
+                                               SourceLocation *MacroEnd) const {
+  assert(Loc.isValid() && Loc.isMacroID() && "Expected a valid macro loc");
+
+  FileID FID = getFileID(Loc);
+  SourceLocation NextLoc = Loc.getLocWithOffset(1);
+  if (isInFileID(NextLoc, FID))
+    return false; // Does not point at the end of expansion range.
+
+  bool Invalid = false;
+  const SrcMgr::ExpansionInfo &ExpInfo =
+      getSLocEntry(FID, &Invalid).getExpansion();
+  if (Invalid)
+    return false;
+
+  if (ExpInfo.isMacroArgExpansion()) {
+    // For macro argument expansions, check if the next FileID is part of the
+    // same argument expansion, in which case this Loc is not at the end of the
+    // expansion.
+    FileID NextFID = getNextFileID(FID);
+    if (!NextFID.isInvalid()) {
+      const SrcMgr::SLocEntry &NextEntry = getSLocEntry(NextFID, &Invalid);
+      if (Invalid)
+        return false;
+      if (NextEntry.isExpansion() &&
+          NextEntry.getExpansion().getExpansionLocStart() ==
+              ExpInfo.getExpansionLocStart())
+        return false;
+    }
+  }
+
+  if (MacroEnd)
+    *MacroEnd = ExpInfo.getExpansionLocEnd();
+  return true;
+}
+
 
 //===----------------------------------------------------------------------===//
 // Queries about the code at a SourceLocation.
diff --git a/lib/Lex/Lexer.cpp b/lib/Lex/Lexer.cpp
index 66498b1..e58581e 100644
--- a/lib/Lex/Lexer.cpp
+++ b/lib/Lex/Lexer.cpp
@@ -798,14 +798,10 @@
                                       SourceLocation *MacroBegin) {
   assert(loc.isValid() && loc.isMacroID() && "Expected a valid macro loc");
 
-  std::pair<FileID, unsigned> infoLoc = SM.getDecomposedLoc(loc);
-  // FIXME: If the token comes from the macro token paste operator ('##')
-  // this function will always return false;
-  if (infoLoc.second > 0)
-    return false; // Does not point at the start of token.
+  SourceLocation expansionLoc;
+  if (!SM.isAtStartOfImmediateMacroExpansion(loc, &expansionLoc))
+    return false;
 
-  SourceLocation expansionLoc =
-    SM.getSLocEntry(infoLoc.first).getExpansion().getExpansionLocStart();
   if (expansionLoc.isFileID()) {
     // No other macro expansions, this is the first.
     if (MacroBegin)
@@ -829,16 +825,11 @@
   if (tokLen == 0)
     return false;
 
-  FileID FID = SM.getFileID(loc);
-  SourceLocation afterLoc = loc.getLocWithOffset(tokLen+1);
-  if (SM.isInFileID(afterLoc, FID))
-    return false; // Still in the same FileID, does not point to the last token.
+  SourceLocation afterLoc = loc.getLocWithOffset(tokLen);
+  SourceLocation expansionLoc;
+  if (!SM.isAtEndOfImmediateMacroExpansion(afterLoc, &expansionLoc))
+    return false;
 
-  // FIXME: If the token comes from the macro token paste operator ('##')
-  // or the stringify operator ('#') this function will always return false;
-
-  SourceLocation expansionLoc =
-    SM.getSLocEntry(FID).getExpansion().getExpansionLocEnd();
   if (expansionLoc.isFileID()) {
     // No other macro expansions.
     if (MacroEnd)
@@ -916,25 +907,25 @@
     return makeRangeFromFileLocs(Range, SM, LangOpts);
   }
 
-  FileID FID;
-  unsigned BeginOffs;
-  llvm::tie(FID, BeginOffs) = SM.getDecomposedLoc(Begin);
-  if (FID.isInvalid())
+  bool Invalid = false;
+  const SrcMgr::SLocEntry &BeginEntry = SM.getSLocEntry(SM.getFileID(Begin),
+                                                        &Invalid);
+  if (Invalid)
     return CharSourceRange();
 
-  unsigned EndOffs;
-  if (!SM.isInFileID(End, FID, &EndOffs) ||
-      BeginOffs > EndOffs)
-    return CharSourceRange();
+  if (BeginEntry.getExpansion().isMacroArgExpansion()) {
+    const SrcMgr::SLocEntry &EndEntry = SM.getSLocEntry(SM.getFileID(End),
+                                                        &Invalid);
+    if (Invalid)
+      return CharSourceRange();
 
-  const SrcMgr::SLocEntry *E = &SM.getSLocEntry(FID);
-  const SrcMgr::ExpansionInfo &Expansion = E->getExpansion();
-  if (Expansion.isMacroArgExpansion() &&
-      Expansion.getSpellingLoc().isFileID()) {
-    SourceLocation SpellLoc = Expansion.getSpellingLoc();
-    Range.setBegin(SpellLoc.getLocWithOffset(BeginOffs));
-    Range.setEnd(SpellLoc.getLocWithOffset(EndOffs));
-    return makeRangeFromFileLocs(Range, SM, LangOpts);
+    if (EndEntry.getExpansion().isMacroArgExpansion() &&
+        BeginEntry.getExpansion().getExpansionLocStart() ==
+            EndEntry.getExpansion().getExpansionLocStart()) {
+      Range.setBegin(SM.getImmediateSpellingLoc(Begin));
+      Range.setEnd(SM.getImmediateSpellingLoc(End));
+      return makeFileCharRange(Range, SM, LangOpts);
+    }
   }
 
   return CharSourceRange();