Fix fn placement around keywords (#613)

Co-authored-by: Geoff Romer <gromer@google.com>

migrate_cpp/cpp_refactoring/fn_inserter.cpp

@@ -5,6 +5,7 @@
 #include "migrate_cpp/cpp_refactoring/fn_inserter.h"
 
 #include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Lex/Lexer.h"
 
 namespace cam = ::clang::ast_matchers;
 
@@ -27,10 +28,29 @@ void FnInserter::run(const cam::MatchFinder::MatchResult& result) {
   if (!decl) {
     llvm::report_fatal_error(std::string("getNodeAs failed for ") + Label);
   }
-  auto begin = decl->getBeginLoc();
-  // Replace the first token in the range, `auto`.
-  auto range = clang::CharSourceRange::getTokenRange(begin, begin);
-  AddReplacement(*(result.SourceManager), range, "fn");
+
+  auto& sm = *(result.SourceManager);
+  auto lang_opts = result.Context->getLangOpts();
+
+  // For names like "Class::Method", replace up to "Class" not "Method".
+  clang::NestedNameSpecifierLoc qual_loc = decl->getQualifierLoc();
+  clang::SourceLocation name_begin_loc =
+      qual_loc.hasQualifier() ? qual_loc.getBeginLoc() : decl->getLocation();
+  auto range =
+      clang::CharSourceRange::getCharRange(decl->getBeginLoc(), name_begin_loc);
+
+  // In order to handle keywords like "virtual" in "virtual auto Foo() -> ...",
+  // scan the replaced text and only drop auto/void entries.
+  llvm::SmallVector<llvm::StringRef> split;
+  clang::Lexer::getSourceText(range, sm, lang_opts)
+      .split(split, ' ', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+  std::string new_text = "fn ";
+  for (llvm::StringRef t : split) {
+    if (t != "auto" && t != "void") {
+      new_text += t.str() + " ";
+    }
+  }
+  AddReplacement(*(result.SourceManager), range, new_text);
 }
 
 }  // namespace Carbon

migrate_cpp/cpp_refactoring/fn_inserter_test.cpp

@@ -23,9 +23,8 @@ TEST_F(FnInserterTest, TrailingReturn) {
 }
 
 TEST_F(FnInserterTest, Inline) {
-  // TODO: Need to re-lex tokens, this should probably be "fn inline" for now.
   constexpr char Before[] = "inline auto A() -> int;";
-  constexpr char After[] = "fn auto A() -> int;";
+  constexpr char After[] = "fn inline A() -> int;";
   ExpectReplacement(Before, After);
 }
 
@@ -36,7 +35,6 @@ TEST_F(FnInserterTest, Void) {
 }
 
 TEST_F(FnInserterTest, Methods) {
-  // TODO: Need to re-lex tokens, this should probably be "fn virtual" for now.
   constexpr char Before[] = R"cpp(
     class Shape {
      public:
@@ -53,12 +51,14 @@ TEST_F(FnInserterTest, Methods) {
      private:
       double radius_;
     };
+
+    void Shape::Draw() {}
   )cpp";
   constexpr char After[] = R"(
     class Shape {
      public:
-      fn void Draw() = 0;
-      fn auto NumSides() -> int = 0;
+      fn virtual Draw() = 0;
+      fn virtual NumSides() -> int = 0;
     };
 
     class Circle : public Shape {
@@ -70,6 +70,8 @@ TEST_F(FnInserterTest, Methods) {
      private:
       double radius_;
     };
+
+    fn Shape::Draw() {}
   )";
   ExpectReplacement(Before, After);
 }