[Refactoring] Replace lifted breaks/returns with placeholders

If we're lifting them outside of the control flow
structure they're dealing with, turn them into
placeholders, as they will no longer perform the
control flow the user is expecting.

This handles:
- Return statements at the top-level of the callback.
- Break statements in switches that we re-write.

Resolves rdar://74014897.

Author: hamishknight

lib/IDE/Refactoring.cpp

@@ -4517,6 +4517,7 @@ struct CallbackClassifier {
   /// names from `Body`. Errors are added through `DiagEngine`, possibly
   /// resulting in partially filled out blocks.
   static void classifyInto(ClassifiedBlocks &Blocks,
+                           llvm::DenseSet<SwitchStmt *> &HandledSwitches,
                            DiagnosticEngine &DiagEngine,
                            ArrayRef<const ParamDecl *> SuccessParams,
                            const ParamDecl *ErrParam, HandlerType ResultType,
@@ -4528,25 +4529,30 @@ struct CallbackClassifier {
     if (ErrParam)
       ParamsSet.insert(ErrParam);
 
-    CallbackClassifier Classifier(Blocks, DiagEngine, ParamsSet, ErrParam,
+    CallbackClassifier Classifier(Blocks, HandledSwitches, DiagEngine,
+                                  ParamsSet, ErrParam,
                                   ResultType == HandlerType::RESULT);
     Classifier.classifyNodes(Body);
   }
 
 private:
   ClassifiedBlocks &Blocks;
+  llvm::DenseSet<SwitchStmt *> &HandledSwitches;
   DiagnosticEngine &DiagEngine;
   ClassifiedBlock *CurrentBlock;
   llvm::DenseSet<const Decl *> ParamsSet;
   const ParamDecl *ErrParam;
   bool IsResultParam;
 
-  CallbackClassifier(ClassifiedBlocks &Blocks, DiagnosticEngine &DiagEngine,
+  CallbackClassifier(ClassifiedBlocks &Blocks,
+                     llvm::DenseSet<SwitchStmt *> &HandledSwitches,
+                     DiagnosticEngine &DiagEngine,
                      llvm::DenseSet<const Decl *> ParamsSet,
                      const ParamDecl *ErrParam, bool IsResultParam)
-      : Blocks(Blocks), DiagEngine(DiagEngine),
-        CurrentBlock(&Blocks.SuccessBlock), ParamsSet(ParamsSet),
-        ErrParam(ErrParam), IsResultParam(IsResultParam) {}
+      : Blocks(Blocks), HandledSwitches(HandledSwitches),
+        DiagEngine(DiagEngine), CurrentBlock(&Blocks.SuccessBlock),
+        ParamsSet(ParamsSet), ErrParam(ErrParam), IsResultParam(IsResultParam) {
+  }
 
   void classifyNodes(ArrayRef<ASTNode> Nodes) {
     for (auto I = Nodes.begin(), E = Nodes.end(); I < E; ++I) {
@@ -4718,6 +4724,8 @@ struct CallbackClassifier {
       if (DiagEngine.hadAnyError())
         return;
     }
+    // Mark this switch statement as having been transformed.
+    HandledSwitches.insert(SS);
   }
 };
 
@@ -4776,6 +4784,9 @@ class AsyncConverter : private SourceEntityWalker {
   // references to it
   llvm::DenseMap<const Decl *, std::string> Names;
 
+  /// The switch statements that have been re-written by this transform.
+  llvm::DenseSet<SwitchStmt *> HandledSwitches;
+
   // These are per-node (ie. are saved and restored on each convertNode call)
   SourceLoc LastAddedLoc;
   int NestedExprCount = 0;
@@ -4846,6 +4857,9 @@ class AsyncConverter : private SourceEntityWalker {
       NestedExprCount++;
       return true;
     }
+    // Note we don't walk into any nested local function decls. If we start
+    // doing so in the future, be sure to update the logic that deals with
+    // converting unhandled returns into placeholders in walkToStmtPre.
     return false;
   }
 
@@ -4905,6 +4919,38 @@ class AsyncConverter : private SourceEntityWalker {
     NestedExprCount++;
     return true;
   }
+
+  bool replaceRangeWithPlaceholder(SourceRange range) {
+    return addCustom(range, [&]() {
+      OS << PLACEHOLDER_START;
+      addRange(range, /*toEndOfToken*/ true);
+      OS << PLACEHOLDER_END;
+    });
+  }
+
+  bool walkToStmtPre(Stmt *S) override {
+    // Some break and return statements need to be turned into placeholders,
+    // as they may no longer perform the control flow that the user is
+    // expecting.
+    if (!S->isImplicit()) {
+      // For a break, if it's jumping out of a switch statement that we've
+      // re-written as a part of the transform, turn it into a placeholder, as
+      // it would have been lifted out of the switch statement.
+      if (auto *BS = dyn_cast<BreakStmt>(S)) {
+        if (auto *SS = dyn_cast<SwitchStmt>(BS->getTarget())) {
+          if (HandledSwitches.contains(SS))
+            replaceRangeWithPlaceholder(S->getSourceRange());
+        }
+      }
+
+      // For a return, if it's not nested inside another closure or function,
+      // turn it into a placeholder, as it will be lifted out of the callback.
+      if (isa<ReturnStmt>(S) && NestedExprCount == 0)
+        replaceRangeWithPlaceholder(S->getSourceRange());
+    }
+    return true;
+  }
+
 #undef PLACEHOLDER_START
 #undef PLACEHOLDER_END
 
@@ -5107,9 +5153,9 @@ class AsyncConverter : private SourceEntityWalker {
     if (!HandlerDesc.HasError) {
       Blocks.SuccessBlock.addAllNodes(CallbackBody);
     } else if (!CallbackBody.empty()) {
-      CallbackClassifier::classifyInto(Blocks, DiagEngine, SuccessParams,
-                                       ErrParam, HandlerDesc.Type,
-                                       CallbackBody);
+      CallbackClassifier::classifyInto(Blocks, HandledSwitches, DiagEngine,
+                                       SuccessParams, ErrParam,
+                                       HandlerDesc.Type, CallbackBody);
       if (DiagEngine.hadAnyError()) {
         // Can only fallback when the results are params, in which case only
         // the names are used (defaulted to the names of the params if none)

test/refactoring/ConvertAsync/convert_params_single.swift

@@ -357,7 +357,7 @@ withError { res, err in
 // NESTEDRET-NEXT: let str = try await withError()
 // NESTEDRET-NEXT: print("before")
 // NESTEDRET-NEXT: if test(str) {
-// NESTEDRET-NEXT: return
+// NESTEDRET-NEXT:   <#return#>
 // NESTEDRET-NEXT: }
 // NESTEDRET-NEXT: print("got result \(str)")
 // NESTEDRET-NEXT: print("after")

test/refactoring/ConvertAsync/convert_result.swift

@@ -1,4 +1,5 @@
 func simple(_ completion: (Result<String, Error>) -> Void) { }
+func simpleWithArg(_ arg: Int, _ completion: (Result<String, Error>) -> Void) { }
 func noError(_ completion: (Result<String, Never>) -> Void) { }
 func test(_ str: String) -> Bool { return false }
 
@@ -279,7 +280,7 @@ simple { res in
 // NESTEDRET-NEXT: let str = try await simple()
 // NESTEDRET-NEXT: print("before")
 // NESTEDRET-NEXT: if test(str) {
-// NESTEDRET-NEXT: return
+// NESTEDRET-NEXT:   <#return#>
 // NESTEDRET-NEXT: }
 // NESTEDRET-NEXT: print("result \(str)")
 // NESTEDRET-NEXT: print("after")
@@ -303,7 +304,7 @@ simple { res in
 // NESTEDBREAK-NEXT: let str = try await simple()
 // NESTEDBREAK-NEXT: print("before")
 // NESTEDBREAK-NEXT: if test(str) {
-// NESTEDBREAK-NEXT: break
+// NESTEDBREAK-NEXT:   <#break#>
 // NESTEDBREAK-NEXT: }
 // NESTEDBREAK-NEXT: print("result \(str)")
 // NESTEDBREAK-NEXT: print("after")
@@ -344,3 +345,70 @@ simple { res in
 // IGNORE-UNRELATED-NEXT:  {{^}} break{{$}}
 // IGNORE-UNRELATED-NEXT:  }
 // IGNORE-UNRELATED-NEXT: print("after")
+
+// RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=BREAK-RET-PLACEHOLDER %s
+simpleWithArg({ return 0 }()) { res in
+  switch res {
+  case .success:
+    if .random() { break }
+    x: if .random() { break x }
+  case .failure:
+    break
+  }
+
+  func foo<T>(_ x: T) {
+    if .random() { return }
+  }
+  foo(res)
+
+  let fn = {
+    if .random() { return }
+    return
+  }
+  fn()
+
+  _ = { return }()
+
+  switch Bool.random() {
+  case true:
+    break
+  case false:
+    if .random() { break }
+    y: if .random() { break y }
+    return
+  }
+
+  x: if .random() {
+    break x
+  }
+  if .random() { return }
+}
+
+// Make sure we replace lifted break/returns with placeholders, but keep nested
+// break/returns in e.g closures or labelled control flow in place.
+
+// BREAK-RET-PLACEHOLDER:      let res = try await simpleWithArg({ return 0 }())
+// BREAK-RET-PLACEHOLDER-NEXT: if .random() { <#break#> }
+// BREAK-RET-PLACEHOLDER-NEXT: x: if .random() { break x }
+// BREAK-RET-PLACEHOLDER-NEXT: func foo<T>(_ x: T) {
+// BREAK-RET-PLACEHOLDER-NEXT:   if .random() { return }
+// BREAK-RET-PLACEHOLDER-NEXT: }
+// BREAK-RET-PLACEHOLDER-NEXT: foo(<#res#>)
+// BREAK-RET-PLACEHOLDER-NEXT: let fn = {
+// BREAK-RET-PLACEHOLDER-NEXT:   if .random() { return }
+// BREAK-RET-PLACEHOLDER-NEXT:   {{^}} return{{$}}
+// BREAK-RET-PLACEHOLDER-NEXT: }
+// BREAK-RET-PLACEHOLDER-NEXT: fn()
+// BREAK-RET-PLACEHOLDER-NEXT: _ = { return }()
+// BREAK-RET-PLACEHOLDER-NEXT: switch Bool.random() {
+// BREAK-RET-PLACEHOLDER-NEXT: case true:
+// BREAK-RET-PLACEHOLDER-NEXT:   {{^}} break{{$}}
+// BREAK-RET-PLACEHOLDER-NEXT: case false:
+// BREAK-RET-PLACEHOLDER-NEXT:   if .random() { break }
+// BREAK-RET-PLACEHOLDER-NEXT:   y: if .random() { break y }
+// BREAK-RET-PLACEHOLDER-NEXT:   <#return#>
+// BREAK-RET-PLACEHOLDER-NEXT: }
+// BREAK-RET-PLACEHOLDER-NEXT: x: if .random() {
+// BREAK-RET-PLACEHOLDER-NEXT:   {{^}} break x{{$}}
+// BREAK-RET-PLACEHOLDER-NEXT: }
+// BREAK-RET-PLACEHOLDER-NEXT: if .random() { <#return#> }