[refactor][extract] insert semicolons into extracted/inserted code

when needed

This commit implements the semicolon insertion logic into the extract
refactoring. The following rules are used:

- extracting expression: add terminating ';' to the extracted function.
- extracting statements that don't require terminating ';' (e.g. switch): add
  terminating ';' to the callee.
- extracting statements with ';':  move (if possible) the original ';' from the
  callee and add terminating ';'.
- otherwise, add ';' to both places.

Differential Revision: https://reviews.llvm.org/D39441


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

4 files changed, 173 insertions, 4 deletions

diff --git a/lib/Tooling/Refactoring/CMakeLists.txt b/lib/Tooling/Refactoring/CMakeLists.txt
index 5d3ddd45b6..402b5d3c6a 100644
--- a/lib/Tooling/Refactoring/CMakeLists.txt
+++ b/lib/Tooling/Refactoring/CMakeLists.txt
@@ -4,7 +4,8 @@ add_clang_library(clangToolingRefactor
   ASTSelection.cpp
   ASTSelectionRequirements.cpp
   AtomicChange.cpp
-  Extract.cpp
+  Extract/Extract.cpp
+  Extract/SourceExtraction.cpp
   RefactoringActions.cpp
   Rename/RenamingAction.cpp
   Rename/SymbolOccurrences.cpp
diff --git a/lib/Tooling/Refactoring/Extract.cpp b/lib/Tooling/Refactoring/Extract.cpp
index 3f5a839318..b0847a7400 100644
--- a/lib/Tooling/Refactoring/Extract.cpp
+++ b/lib/Tooling/Refactoring/Extract.cpp
@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Tooling/Refactoring/Extract/Extract.h"
+#include "SourceExtraction.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
@@ -145,6 +146,8 @@ ExtractFunction::createSourceReplacements(RefactoringRuleContext &Context) {
   PP.SuppressLifetimeQualifiers = true;
   PP.SuppressUnwrittenScope = true;
 
+  ExtractionSemicolonPolicy Semicolons = ExtractionSemicolonPolicy::compute(
+      Code[Code.size() - 1], ExtractedRange, SM, LangOpts);
   AtomicChange Change(SM, ExtractedDeclLocation);
   // Create the replacement for the extracted declaration.
   {
@@ -162,8 +165,8 @@ ExtractFunction::createSourceReplacements(RefactoringRuleContext &Context) {
     if (IsExpr && !ReturnType->isVoidType())
       OS << "return ";
     OS << ExtractedCodeRewriter.getRewrittenText(ExtractedRange);
-    // FIXME: Compute the correct semicolon policy.
-    OS << ';';
+    if (Semicolons.isNeededInExtractedFunction())
+      OS << ';';
     OS << "\n}\n\n";
     auto Err = Change.insert(SM, ExtractedDeclLocation, OS.str());
     if (Err)
@@ -178,7 +181,8 @@ ExtractFunction::createSourceReplacements(RefactoringRuleContext &Context) {
     OS << DeclName << '(';
     // FIXME: Forward arguments.
     OS << ')';
-    // FIXME: Add semicolon if needed.
+    if (Semicolons.isNeededInOriginalFunction())
+      OS << ';';
 
     auto Err = Change.replace(
         SM, CharSourceRange::getTokenRange(ExtractedRange), OS.str());
diff --git a/lib/Tooling/Refactoring/Extract/SourceExtraction.cpp b/lib/Tooling/Refactoring/Extract/SourceExtraction.cpp
new file mode 100644
index 0000000000..7fd8cc2d3c
--- /dev/null
+++ b/lib/Tooling/Refactoring/Extract/SourceExtraction.cpp
@@ -0,0 +1,112 @@
+//===--- SourceExtraction.cpp - Clang refactoring library -----------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SourceExtraction.h"
+#include "clang/AST/Stmt.h"
+#include "clang/AST/StmtCXX.h"
+#include "clang/AST/StmtObjC.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Lex/Lexer.h"
+
+using namespace clang;
+
+namespace {
+
+/// Returns true if the token at the given location is a semicolon.
+bool isSemicolonAtLocation(SourceLocation TokenLoc, const SourceManager &SM,
+                           const LangOptions &LangOpts) {
+  return Lexer::getSourceText(
+             CharSourceRange::getTokenRange(TokenLoc, TokenLoc), SM,
+             LangOpts) == ";";
+}
+
+/// Returns true if there should be a semicolon after the given statement.
+bool isSemicolonRequiredAfter(const Stmt *S) {
+  if (isa<CompoundStmt>(S))
+    return false;
+  if (const auto *If = dyn_cast<IfStmt>(S))
+    return isSemicolonRequiredAfter(If->getElse() ? If->getElse()
+                                                  : If->getThen());
+  if (const auto *While = dyn_cast<WhileStmt>(S))
+    return isSemicolonRequiredAfter(While->getBody());
+  if (const auto *For = dyn_cast<ForStmt>(S))
+    return isSemicolonRequiredAfter(For->getBody());
+  if (const auto *CXXFor = dyn_cast<CXXForRangeStmt>(S))
+    return isSemicolonRequiredAfter(CXXFor->getBody());
+  if (const auto *ObjCFor = dyn_cast<ObjCForCollectionStmt>(S))
+    return isSemicolonRequiredAfter(ObjCFor->getBody());
+  switch (S->getStmtClass()) {
+  case Stmt::SwitchStmtClass:
+  case Stmt::CXXTryStmtClass:
+  case Stmt::ObjCAtSynchronizedStmtClass:
+  case Stmt::ObjCAutoreleasePoolStmtClass:
+  case Stmt::ObjCAtTryStmtClass:
+    return false;
+  default:
+    return true;
+  }
+}
+
+/// Returns true if the two source locations are on the same line.
+bool areOnSameLine(SourceLocation Loc1, SourceLocation Loc2,
+                   const SourceManager &SM) {
+  return !Loc1.isMacroID() && !Loc2.isMacroID() &&
+         SM.getSpellingLineNumber(Loc1) == SM.getSpellingLineNumber(Loc2);
+}
+
+} // end anonymous namespace
+
+namespace clang {
+namespace tooling {
+
+ExtractionSemicolonPolicy
+ExtractionSemicolonPolicy::compute(const Stmt *S, SourceRange &ExtractedRange,
+                                   const SourceManager &SM,
+                                   const LangOptions &LangOpts) {
+  auto neededInExtractedFunction = []() {
+    return ExtractionSemicolonPolicy(true, false);
+  };
+  auto neededInOriginalFunction = []() {
+    return ExtractionSemicolonPolicy(false, true);
+  };
+
+  /// The extracted expression should be terminated with a ';'. The call to
+  /// the extracted function will replace this expression, so it won't need
+  /// a terminating ';'.
+  if (isa<Expr>(S))
+    return neededInExtractedFunction();
+
+  /// Some statements don't need to be terminated with ';'. The call to the
+  /// extracted function will be a standalone statement, so it should be
+  /// terminated with a ';'.
+  bool NeedsSemi = isSemicolonRequiredAfter(S);
+  if (!NeedsSemi)
+    return neededInOriginalFunction();
+
+  /// Some statements might end at ';'. The extraction will move that ';', so
+  /// the call to the extracted function should be terminated with a ';'.
+  SourceLocation End = ExtractedRange.getEnd();
+  if (isSemicolonAtLocation(End, SM, LangOpts))
+    return neededInOriginalFunction();
+
+  /// Other statements should generally have a trailing ';'. We can try to find
+  /// it and move it together it with the extracted code.
+  Optional<Token> NextToken = Lexer::findNextToken(End, SM, LangOpts);
+  if (NextToken && NextToken->is(tok::semi) &&
+      areOnSameLine(NextToken->getLocation(), End, SM)) {
+    ExtractedRange.setEnd(NextToken->getLocation());
+    return neededInOriginalFunction();
+  }
+
+  /// Otherwise insert semicolons in both places.
+  return ExtractionSemicolonPolicy(true, true);
+}
+
+} // end namespace tooling
+} // end namespace clang
diff --git a/lib/Tooling/Refactoring/Extract/SourceExtraction.h b/lib/Tooling/Refactoring/Extract/SourceExtraction.h
new file mode 100644
index 0000000000..4b4bd8b477
--- /dev/null
+++ b/lib/Tooling/Refactoring/Extract/SourceExtraction.h
@@ -0,0 +1,52 @@
+//===--- SourceExtraction.cpp - Clang refactoring library -----------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_LIB_TOOLING_REFACTORING_EXTRACT_SOURCE_EXTRACTION_H
+#define LLVM_CLANG_LIB_TOOLING_REFACTORING_EXTRACT_SOURCE_EXTRACTION_H
+
+#include "clang/Basic/LLVM.h"
+
+namespace clang {
+
+class LangOptions;
+class SourceManager;
+class SourceRange;
+class Stmt;
+
+namespace tooling {
+
+/// Determines which semicolons should be inserted during extraction.
+class ExtractionSemicolonPolicy {
+public:
+  bool isNeededInExtractedFunction() const {
+    return IsNeededInExtractedFunction;
+  }
+
+  bool isNeededInOriginalFunction() const { return IsNeededInOriginalFunction; }
+
+  /// Returns the semicolon insertion policy that is needed for extraction of
+  /// the given statement from the given source range.
+  static ExtractionSemicolonPolicy compute(const Stmt *S,
+                                           SourceRange &ExtractedRange,
+                                           const SourceManager &SM,
+                                           const LangOptions &LangOpts);
+
+private:
+  ExtractionSemicolonPolicy(bool IsNeededInExtractedFunction,
+                            bool IsNeededInOriginalFunction)
+      : IsNeededInExtractedFunction(IsNeededInExtractedFunction),
+        IsNeededInOriginalFunction(IsNeededInOriginalFunction) {}
+  bool IsNeededInExtractedFunction;
+  bool IsNeededInOriginalFunction;
+};
+
+} // end namespace tooling
+} // end namespace clang
+
+#endif // LLVM_CLANG_LIB_TOOLING_REFACTORING_EXTRACT_SOURCE_EXTRACTION_H