[Refactoring] Don't drop returns with expressions

Previously we were unconditionally dropping a
return statement if it was the last node, which
could cause us to inadvertently drop the result
expression as well. Instead, only drop bare
'return' statements.

rdar://77789360

Author: hamishknight

lib/IDE/Refactoring.cpp

@@ -4631,16 +4631,24 @@ class NodesToPrint {
            !Nodes.back().isImplicit();
   }
 
-  /// If the last recorded node is an explicit return or break statement, drop
-  /// it from the list.
-  void dropTrailingReturnOrBreak() {
+  /// If the last recorded node is an explicit return or break statement that
+  /// can be safely dropped, drop it from the list.
+  void dropTrailingReturnOrBreakIfPossible() {
     if (!hasTrailingReturnOrBreak())
       return;
 
+    auto *Node = Nodes.back().get<Stmt *>();
+
+    // If this is a return statement with return expression, let's preserve it.
+    if (auto *RS = dyn_cast<ReturnStmt>(Node)) {
+      if (RS->hasResult())
+        return;
+    }
+
     // Remove the node from the list, but make sure to add it as a possible
     // comment loc to preserve any of its attached comments.
-    auto Node = Nodes.pop_back_val();
-    addPossibleCommentLoc(Node.getStartLoc());
+    Nodes.pop_back();
+    addPossibleCommentLoc(Node->getStartLoc());
   }
 
   /// Returns a list of nodes to print in a brace statement. This picks up the
@@ -4903,7 +4911,7 @@ struct CallbackClassifier {
                 NodesToPrint Nodes) {
     if (Nodes.hasTrailingReturnOrBreak()) {
       CurrentBlock = OtherBlock;
-      Nodes.dropTrailingReturnOrBreak();
+      Nodes.dropTrailingReturnOrBreakIfPossible();
       Block->addAllNodes(std::move(Nodes));
     } else {
       Block->addAllNodes(std::move(Nodes));

test/refactoring/ConvertAsync/convert_function.swift

@@ -257,3 +257,26 @@ func testReturnHandling(_ completion: (String?, Error?) -> Void) {
 // RETURN-HANDLING-NEXT:   {{^}} return ""{{$}}
 // RETURN-HANDLING-NEXT: }
 
+// rdar://77789360 Make sure we don't print a double return statement and don't
+// completely drop completion(a).
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=RETURN-HANDLING2 %s
+func testReturnHandling2(completion: @escaping (String) -> ()) {
+  testReturnHandling { x, err in
+    guard let x = x else {
+      let a = ""
+      return completion(a)
+    }
+    let b = ""
+    return completion(b)
+  }
+}
+// RETURN-HANDLING2:      func testReturnHandling2() async -> String {
+// RETURN-HANDLING2-NEXT:   do {
+// RETURN-HANDLING2-NEXT:     let x = try await testReturnHandling()
+// RETURN-HANDLING2-NEXT:     let b = ""
+// RETURN-HANDLING2-NEXT:     {{^}}<#return#> b{{$}}
+// RETURN-HANDLING2-NEXT:   } catch let err {
+// RETURN-HANDLING2-NEXT:     let a = ""
+// RETURN-HANDLING2-NEXT:     {{^}}<#return#> a{{$}}
+// RETURN-HANDLING2-NEXT:   }
+// RETURN-HANDLING2-NEXT: }