[clang-move] Fix incorrect EndLoc for declarations in macros.

[ioeric]

Lg

clang-move/ClangMove.cpp

 };
 
 // Expand to get the end location of the line where the EndLoc of the given
 // Decl.
 SourceLocation
 getLocForEndOfDecl(const clang::Decl *D,
                    const LangOptions &LangOpts = clang::LangOptions()) {
   const auto &SM = D->getASTContext().getSourceManager();
-  std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(D->getLocEnd());
+  auto EndExpansionLoc = SM.getExpansionLoc(D->getLocEnd());
+  std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(EndExpansionLoc);
   // Try to load the file buffer.
   bool InvalidTemp = false;
   llvm::StringRef File = SM.getBufferData(LocInfo.first, &InvalidTemp);
   if (InvalidTemp)
     return SourceLocation();
 
   const char *TokBegin = File.data() + LocInfo.second;
   // Lex from the start of the given location.
   Lexer Lex(SM.getLocForStartOfFile(LocInfo.first), LangOpts, File.begin(),
             TokBegin, File.end());
 
   llvm::SmallVector<char, 16> Line;
   // FIXME: this is a bit hacky to get ReadToEndOfLine work.
   Lex.setParsingPreprocessorDirective(true);
   Lex.ReadToEndOfLine(&Line);
-  SourceLocation EndLoc = D->getLocEnd().getLocWithOffset(Line.size());
+  SourceLocation EndLoc = EndExpansionLoc.getLocWithOffset(Line.size());
   // If we already reach EOF, just return the EOF SourceLocation;
   // otherwise, move 1 offset ahead to include the trailing newline character
   // '\n'.
   return SM.getLocForEndOfFile(LocInfo.first) == EndLoc
              ? EndLoc
              : EndLoc.getLocWithOffset(1);
 }
 

unittests/clang-move/ClangMoveTests.cpp

TEST(ClangMove, MacroInFunction) {
   Spec.OldHeader = "foo.h";
   Spec.OldCC = "foo.cc";
   Spec.NewHeader = "new_foo.h";
   Spec.NewCC = "new_foo.cc";
   auto Results = runClangMoveOnCode(Spec, TestHeader, TestCode);
   EXPECT_EQ(ExpectedNewCode, Results[Spec.NewCC]);
 }
 
+TEST(ClangMove, DefinitionInMacro) {
+  const char TestHeader[] = "#define DEF(CLASS) void CLASS##_::f() {}\n"
+                            "class A_ {\nvoid f();\n};\n"
+                            "class B {};\n";
+  const char TestCode[] = "#include \"foo.h\"\n"
+                          "DEF(A)\n";
+  const char ExpectedNewCode[] = "#include \"new_foo.h\"\n\n"
+                                 "DEF(A)\n";
+  move::MoveDefinitionSpec Spec;
+  Spec.Names.push_back("A_");
+  Spec.OldHeader = "foo.h";
+  Spec.OldCC = "foo.cc";
+  Spec.NewHeader = "new_foo.h";
+  Spec.NewCC = "new_foo.cc";
+  auto Results = runClangMoveOnCode(Spec, TestHeader, TestCode);
+  EXPECT_EQ(ExpectedNewCode, Results[Spec.NewCC]);
+}
+
 TEST(ClangMove, WellFormattedCode) {
   const std::string CommonHeader =
       "namespace a {\n"
       "namespace b {\n"
       "namespace c {\n"
       "class C;\n"
       "\n"
       "class A {\npublic:\n  void f();\n  void f2();\n};\n"