1 files changed, 68 insertions, 24 deletions

diff --git a/clang/lib/Tooling/Transformer/SourceCode.cpp b/clang/lib/Tooling/Transformer/SourceCode.cpp
index 26b204851f05..35edc261ef09 100644
--- a/clang/lib/Tooling/Transformer/SourceCode.cpp
+++ b/clang/lib/Tooling/Transformer/SourceCode.cpp
@@ -50,8 +50,9 @@ CharSourceRange clang::tooling::maybeExtendRange(CharSourceRange Range,
   return CharSourceRange::getTokenRange(Range.getBegin(), Tok.getLocation());
 }
 
-llvm::Error clang::tooling::validateEditRange(const CharSourceRange &Range,
-                                              const SourceManager &SM) {
+static llvm::Error validateRange(const CharSourceRange &Range,
+                                 const SourceManager &SM,
+                                 bool AllowSystemHeaders) {
   if (Range.isInvalid())
     return llvm::make_error<StringError>(errc::invalid_argument,
                                          "Invalid range");
@@ -60,10 +61,12 @@ llvm::Error clang::tooling::validateEditRange(const CharSourceRange &Range,
     return llvm::make_error<StringError>(
         errc::invalid_argument, "Range starts or ends in a macro expansion");
 
-  if (SM.isInSystemHeader(Range.getBegin()) ||
-      SM.isInSystemHeader(Range.getEnd()))
-    return llvm::make_error<StringError>(errc::invalid_argument,
-                                         "Range is in system header");
+  if (!AllowSystemHeaders) {
+    if (SM.isInSystemHeader(Range.getBegin()) ||
+        SM.isInSystemHeader(Range.getEnd()))
+      return llvm::make_error<StringError>(errc::invalid_argument,
+                                           "Range is in system header");
+  }
 
   std::pair<FileID, unsigned> BeginInfo = SM.getDecomposedLoc(Range.getBegin());
   std::pair<FileID, unsigned> EndInfo = SM.getDecomposedLoc(Range.getEnd());
@@ -72,33 +75,74 @@ llvm::Error clang::tooling::validateEditRange(const CharSourceRange &Range,
         errc::invalid_argument, "Range begins and ends in different files");
 
   if (BeginInfo.second > EndInfo.second)
-    return llvm::make_error<StringError>(
-        errc::invalid_argument, "Range's begin is past its end");
+    return llvm::make_error<StringError>(errc::invalid_argument,
+                                         "Range's begin is past its end");
 
   return llvm::Error::success();
 }
 
-llvm::Optional<CharSourceRange>
-clang::tooling::getRangeForEdit(const CharSourceRange &EditRange,
+llvm::Error clang::tooling::validateEditRange(const CharSourceRange &Range,
+                                              const SourceManager &SM) {
+  return validateRange(Range, SM, /*AllowSystemHeaders=*/false);
+}
+
+static bool spelledInMacroDefinition(SourceLocation Loc,
+                                     const SourceManager &SM) {
+  while (Loc.isMacroID()) {
+    const auto &Expansion = SM.getSLocEntry(SM.getFileID(Loc)).getExpansion();
+    if (Expansion.isMacroArgExpansion()) {
+      // Check the spelling location of the macro arg, in case the arg itself is
+      // in a macro expansion.
+      Loc = Expansion.getSpellingLoc();
+    } else {
+      return true;
+    }
+  }
+  return false;
+}
+
+static CharSourceRange getRange(const CharSourceRange &EditRange,
                                 const SourceManager &SM,
-                                const LangOptions &LangOpts) {
-  // FIXME: makeFileCharRange() has the disadvantage of stripping off "identity"
-  // macros. For example, if we're looking to rewrite the int literal 3 to 6,
-  // and we have the following definition:
-  //    #define DO_NOTHING(x) x
-  // then
-  //    foo(DO_NOTHING(3))
-  // will be rewritten to
-  //    foo(6)
-  // rather than the arguably better
-  //    foo(DO_NOTHING(6))
-  // Decide whether the current behavior is desirable and modify if not.
-  CharSourceRange Range = Lexer::makeFileCharRange(EditRange, SM, LangOpts);
+                                const LangOptions &LangOpts,
+                                bool IncludeMacroExpansion) {
+  CharSourceRange Range;
+  if (IncludeMacroExpansion) {
+    Range = Lexer::makeFileCharRange(EditRange, SM, LangOpts);
+  } else {
+    if (spelledInMacroDefinition(EditRange.getBegin(), SM) ||
+        spelledInMacroDefinition(EditRange.getEnd(), SM))
+      return {};
+
+    auto B = SM.getSpellingLoc(EditRange.getBegin());
+    auto E = SM.getSpellingLoc(EditRange.getEnd());
+    if (EditRange.isTokenRange())
+      E = Lexer::getLocForEndOfToken(E, 0, SM, LangOpts);
+    Range = CharSourceRange::getCharRange(B, E);
+  }
+  return Range;
+}
+
+std::optional<CharSourceRange> clang::tooling::getFileRangeForEdit(
+    const CharSourceRange &EditRange, const SourceManager &SM,
+    const LangOptions &LangOpts, bool IncludeMacroExpansion) {
+  CharSourceRange Range =
+      getRange(EditRange, SM, LangOpts, IncludeMacroExpansion);
   bool IsInvalid = llvm::errorToBool(validateEditRange(Range, SM));
   if (IsInvalid)
-    return llvm::None;
+    return std::nullopt;
   return Range;
+}
 
+std::optional<CharSourceRange> clang::tooling::getFileRange(
+    const CharSourceRange &EditRange, const SourceManager &SM,
+    const LangOptions &LangOpts, bool IncludeMacroExpansion) {
+  CharSourceRange Range =
+      getRange(EditRange, SM, LangOpts, IncludeMacroExpansion);
+  bool IsInvalid =
+      llvm::errorToBool(validateRange(Range, SM, /*AllowSystemHeaders=*/true));
+  if (IsInvalid)
+    return std::nullopt;
+  return Range;
 }
 
 static bool startsWithNewline(const SourceManager &SM, const Token &Tok) {