Emit fix-it for diagnostic

This patch teaches the compiler to emit a fix-it to automatically wrap
everything in the function in a Task and remove the effects from the
function declaration.

Author: etcwilde

lib/Sema/TypeCheckAttr.cpp

@@ -475,6 +475,68 @@ void AttributeChecker::visitDynamicAttr(DynamicAttr *attr) {
     diagnoseAndRemoveAttr(attr, diag::dynamic_with_transparent);
 }
 
+/// Replaces asynchronous IBActionAttr/IBSegueActionAttr function declarations
+/// with a synchronous function. The body of the original function is moved
+/// inside of a task executed on the MainActor
+static void emitFixItIBActionRemoveAsync(ASTContext &ctx, const FuncDecl &FD) {
+  // If we don't have an async loc for some reason, things will explode
+  if (!FD.getAsyncLoc())
+    return;
+
+  std::string replacement = "";
+
+  // attributes, function name and everything up to `async` (exclusive)
+  replacement += CharSourceRange(ctx.SourceMgr,
+      FD.getSourceRangeIncludingAttrs().Start,
+      FD.getAsyncLoc()).str();
+
+  CharSourceRange returnType = Lexer::getCharSourceRangeFromSourceRange(ctx.SourceMgr,
+      FD.getResultTypeSourceRange());
+
+
+  // If we have a return type, include that here
+  if (returnType.isValid()) {
+    replacement += (llvm::Twine("-> ") + Lexer::getCharSourceRangeFromSourceRange(ctx.SourceMgr,
+          FD.getResultTypeSourceRange()).str()).str();
+  }
+
+
+  if (!FD.hasBody()) {
+    // If we don't have any body, the sourcelocs won't work and will result in
+    // crashes, so just swap out what we can
+
+    SourceLoc endLoc = returnType.isValid() ? returnType.getEnd() : FD.getAsyncLoc();
+    ctx.Diags.diagnose(FD.getAsyncLoc(), diag::remove_async_add_task,
+        FD.getName())
+      .fixItReplace(SourceRange(FD.getSourceRangeIncludingAttrs().Start, endLoc),
+          replacement);
+    return;
+  }
+
+  if (returnType.isValid())
+    replacement += " "; // insert space between type name and lbrace
+
+  replacement += "{\nTask { @MainActor in";
+
+  // If the body of the function is just "{}", there isn't anything to wrap.
+  // stepping over the braces to grab just the body will result in the `Start`
+  // location of the source range to come after the `End` of the range, and we
+  // will overflow. Dance around this by just appending the end of the fix to
+  // the replacement.
+  if (FD.getBody()->getLBraceLoc() != FD.getBody()->getRBraceLoc().getAdvancedLocOrInvalid(-1)) {
+    // We actually have a body, so add that to the string
+    CharSourceRange functionBody(ctx.SourceMgr,
+        FD.getBody()->getLBraceLoc().getAdvancedLocOrInvalid(1),
+        FD.getBody()->getRBraceLoc().getAdvancedLocOrInvalid(-1));
+    replacement += functionBody.str();
+  }
+  replacement += " }\n}";
+
+  ctx.Diags.diagnose(FD.getAsyncLoc(), diag::remove_async_add_task, FD.getName())
+    .fixItReplace(SourceRange(FD.getSourceRangeIncludingAttrs().Start,
+          FD.getBody()->getRBraceLoc()), replacement);
+}
+
 static bool
 validateIBActionSignature(ASTContext &ctx, DeclAttribute *attr,
                           const FuncDecl *FD, unsigned minParameters,
@@ -504,9 +566,7 @@ validateIBActionSignature(ASTContext &ctx, DeclAttribute *attr,
   if (FD->isAsyncContext()) {
     ctx.Diags.diagnose(FD->getAsyncLoc(), diag::attr_decl_async,
         attr->getAttrName(), FD->getDescriptiveKind());
-
-    ctx.Diags.diagnose(FD->getAsyncLoc(), diag::remove_async_add_task,
-                       FD->getName());
+    emitFixItIBActionRemoveAsync(ctx, *FD);
     valid = false;
   }
 

test/attr/attr_ibaction.swift

@@ -42,11 +42,24 @@ class IBActionWrapperTy {
   @objc @IBAction
   func evenMoreMagic(_: AnyObject) -> () {} // no-warning
 
-  @available(macOS 10.15, *)
+  @available(SwiftStdlib 5.5, *)
   @IBAction
-  func asyncIBAction(_: AnyObject) async -> () {}
+  func asyncIBActionNoSpace(_: AnyObject) async -> () {}
   // expected-error@-1 {{@IBAction instance method cannot be async}}
-  // expected-note@-2 {{remove 'async' and wrap in 'Task' to use concurrency in 'asyncIBAction'}}
+  // expected-note@-2 {{remove 'async' and wrap in 'Task' to use concurrency in 'asyncIBActionNoSpace'}}{{45:3-47:57=@available(SwiftStdlib 5.5, *)\n  @IBAction\n  func asyncIBActionNoSpace(_: AnyObject) -> () {\nTask { @MainActor in \}\n\}}}
+
+  @available(SwiftStdlib 5.5, *)
+  @IBAction
+  func asyncIBActionWithFullBody(_: AnyObject) async {
+      print("Hello World")
+  }
+  // expected-error@-3 {{@IBAction instance method cannot be async}}
+  // expected-note@-4 {{remove 'async' and wrap in 'Task' to use concurrency in 'asyncIBActionWithFullBody'}}{{51:3-55:4=@available(SwiftStdlib 5.5, *)\n  @IBAction\n  func asyncIBActionWithFullBody(_: AnyObject) {\nTask { @MainActor in\n      print("Hello World")\n  \}\n\}}}
+
+  @available(SwiftStdlib 5.5, *) @IBAction func asyncIBActionNoBody(_: AnyObject) async
+  // expected-error@-1 {{expected '{' in body of function declaration}}
+  // expected-error@-2 {{@IBAction instance method cannot be asynchronous}}
+  // expected-note@-3 {{remove 'async' and wrap in 'Task' to use concurrency in 'asyncIBActionNoBody}}{{3-88=@available(SwiftStdlib 5.5, *) @IBAction func asyncIBActionNoBody(_: AnyObject)}}
 
 }
 

test/attr/attr_ibsegueaction.swift

@@ -45,6 +45,11 @@ class IBActionWrapperTy {
   func moreMagic(_: AnyObject, _: AnyObject, _: AnyObject) -> AnyObject {} // no-warning
   @objc @IBSegueAction
   func evenMoreMagic(_: AnyObject, _: AnyObject, _: AnyObject) -> AnyObject {} // no-warning
+
+  @available(SwiftStdlib 5.5, *) @IBSegueAction
+  func process(_: AnyObject, _: AnyObject, _: AnyObject) async -> AnyObject { }
+  // expected-error@-1 {{@IBSegueAction instance method cannot be asynchronous}}
+  // expected-note@-2 {{remove 'async' and wrap in 'Task' to use concurrency in 'process'}}{{49:3-50:80=@available(SwiftStdlib 5.5, *) @IBSegueAction\n  func process(_: AnyObject, _: AnyObject, _: AnyObject) -> AnyObject {\nTask { @MainActor in \}\n\}}}
 }
 
 struct S { }