[Refactor][Extract] Correctly emit the return type for functions to

account for block extractions and copy blocks when extracting without ARC

(cherry picked from commit 0405317)

rdar://35578265

lib/Tooling/Refactor/Extract.cpp

@@ -1852,20 +1852,22 @@ llvm::Expected<RefactoringResult> ExtractOperation::perform(
       OS << "static ";
     else if (!isMethodExtraction())
       OS << (isInHeader(FunctionExtractionLoc, SM) ? "inline " : "static ");
-    ReturnType.print(OS, PP);
-    OS << ' ';
+    std::string QualifiedName;
+    llvm::raw_string_ostream NameOS(QualifiedName);
     if (isMethodExtraction() && IsDefinition)
-      printEnclosingMethodScope(FunctionLikeParentDecl, OS, PP);
-    unsigned NameOffset = OS.str().size();
-    OS << ExtractedName << '(';
+      printEnclosingMethodScope(FunctionLikeParentDecl, NameOS, PP);
+    NameOS << ExtractedName;
+    NameOS << '(';
     bool IsFirst = true;
     for (const auto &Var : CapturedVariables) {
       if (!IsFirst)
-        OS << ", ";
+        NameOS << ", ";
       IsFirst = false;
-      Var.ParameterType.print(OS, PP, /*PlaceHolder=*/Var.getName());
+      Var.ParameterType.print(NameOS, PP, /*PlaceHolder=*/Var.getName());
     }
-    OS << ')';
+    NameOS << ')';
+    ReturnType.print(OS, PP, NameOS.str());
+    unsigned NameOffset = OS.str().find(ExtractedName);
     if (isMethodExtraction() && isEnclosingMethodConst(FunctionLikeParentDecl))
       OS << " const";
     return RefactoringReplacement::AssociatedSymbolLocation(
@@ -1923,7 +1925,18 @@ llvm::Expected<RefactoringResult> ExtractOperation::perform(
   auto Semicolons = computeSemicolonExtractionPolicy(
       ExtractedStmtRange ? *(ExtractedStmtRange->Last) : S, ExtractedTokenRange,
       SM, LangOpts);
+  bool ShouldCopyBlock = false;
+  if (IsExpr && !LangOpts.ObjCAutoRefCount &&
+      ReturnType->isBlockPointerType()) {
+    // We can't return local blocks directly without ARC; they should be copied.
+    // FIXME: This is overly pessimistic, as we only need the copy for local
+    // blocks.
+    ExtractedOS << "[(";
+    ShouldCopyBlock = true;
+  }
   ExtractedOS << SourceRewriter.getRewrittenText(ExtractedTokenRange);
+  if (ShouldCopyBlock)
+    ExtractedOS << ") copy]";
   if (Semicolons.IsNeededInExtractedFunction)
     ExtractedOS << ';';
   if (CanUseReturnForVariablesUsedAfterwards)

test/Refactor/Extract/extract-objc-property.m

@@ -16,14 +16,14 @@ - (void)test {
   self.item;
 // property-end: -1:12
 // property-name-end: -2:12
-// CHECK: "static HasProperty * extracted(HasProperty *object) {\nreturn object.item;\n}\n\n"
+// CHECK: "static HasProperty *extracted(HasProperty *object) {\nreturn object.item;\n}\n\n"
 
 // implicit-name-begin: +2:8
 // implicit-begin: +1:3
   self.implicitProp;
 // implicit-end: -1:20
 // implicit-name-end: -2:20
-// CHECK: "static HasProperty * extracted(HasProperty *object) {\nreturn object.implicitProp;\n}\n\n"
+// CHECK: "static HasProperty *extracted(HasProperty *object) {\nreturn object.implicitProp;\n}\n\n"
 }
 
 // RUN: clang-refactor-test perform -action extract -selected=property -selected=implicit %s -fobjc-arc | FileCheck %s

test/Refactor/Extract/return-block.m

@@ -0,0 +1,24 @@
+// RUN: clang-refactor-test perform -action extract -selected=%s:16:27-18:4 %s -fobjc-arc | FileCheck %s
+// RUN: clang-refactor-test perform -action extract -selected=%s:16:27-18:4 %s | FileCheck --check-prefix=NOARC %s
+@interface I
+
+@end
+
+@implementation I
+
+- (void) doStuff: (int)x block: (void (^)(int))block {
+
+}
+
+- (void)foo {}
+
+- (void)viewDidLoad {
+  [self doStuff: 2 block: ^(int returnCode) {
+    [self foo];
+  }];
+}
+// CHECK: "static void (^extracted(I *object))(int) {\nreturn ^(int returnCode) {\n    [object foo];\n  };\n}\n\n" [[@LINE-5]]:1 -> [[@LINE-5]]:1
+// CHECK: "extracted(self)" [[@LINE-5]]:27 -> [[@LINE-3]]:4
+// NOARC: "static void (^extracted(I *object))(int) {\nreturn [(^(int returnCode) {\n    [object foo];\n  }) copy];\n}\n\n" [[@LINE-7]]:1 -> [[@LINE-7]]:1
+
+@end

test/Refactor/Extract/return-correct-stl-type.cpp

@@ -28,7 +28,7 @@ typedef BasicString<Traits> String;
 
 void returnCharTypeNotUselessValueType() {
 // CHECK1: "static char extracted(const std::String &x) {\nreturn x.value();\n}\n\n" [[@LINE-1]]:1
-// CHECK1: "static const char * extracted(const std::String &x) {\nreturn x.data();\n}\n\n" [[@LINE-2]]:1
+// CHECK1: "static const char *extracted(const std::String &x) {\nreturn x.data();\n}\n\n" [[@LINE-2]]:1
   std::String x;
 // return-char-begin: +1:9
   (void)x.value();

test/Refactor/Extract/return-objc-generic-argument-type.m

@@ -19,22 +19,22 @@ void foo(Array<NSObject *> *objects) {
 // prop-begin: +1:3
   objects.prop;
 // prop-end: -1:15
-// CHECK: "static NSObject * extracted(Array<NSObject *> *objects) {\nreturn objects.prop;\n}\n\n"
+// CHECK: "static NSObject *extracted(Array<NSObject *> *objects) {\nreturn objects.prop;\n}\n\n"
 // imp-prop-begin: +1:3
   objects.get;
 // imp-prop-end: -1:14
-// CHECK: "static NSObject * extracted(Array<NSObject *> *objects) {\nreturn objects.get;\n}\n\n"
+// CHECK: "static NSObject *extracted(Array<NSObject *> *objects) {\nreturn objects.get;\n}\n\n"
 // class-prop-begin: +1:3
   Array.classProp;
 // class-prop-end: -1:30
-// CHECK: "static Array * extracted() {\nreturn Array.classProp;\n}\n\n"
+// CHECK: "static Array *extracted() {\nreturn Array.classProp;\n}\n\n"
   typedef Array<NSObject *> ObjectArray;
 // class-prop2-begin: +1:3
   [ObjectArray classProp];
 // class-prop2-end: -1:26
-// CHECK: "static Array<NSObject *> * extracted() {\nreturn [ObjectArray classProp];\n}\n\n"
+// CHECK: "static Array<NSObject *> *extracted() {\nreturn [ObjectArray classProp];\n}\n\n"
 // class-method-begin: +1:3
   [ObjectArray classGet];
 // class-method-end: -1:25
-// CHECK: "static NSObject ** extracted() {\nreturn [ObjectArray classGet];\n}\n\n"
+// CHECK: "static NSObject **extracted() {\nreturn [ObjectArray classGet];\n}\n\n"
 }