[Async Refactoring] Better handle known condition paths

lib/IDE/Refactoring.cpp

@@ -4933,6 +4933,12 @@ struct CallbackClassifier {
   const ParamDecl *ErrParam;
   bool IsResultParam;
 
+  /// This is set to \c true if we're currently classifying on a known condition
+  /// path, where \c CurrentBlock is set to the appropriate block. This lets us
+  /// be more lenient with unhandled conditions as we already know the block
+  /// we're supposed to be in.
+  bool IsKnownConditionPath = false;
+
   CallbackClassifier(ClassifiedBlocks &Blocks, const FuncDecl *Callee,
                      ArrayRef<const ParamDecl *> SuccessParams,
                      llvm::DenseSet<SwitchStmt *> &HandledSwitches,
@@ -5248,7 +5254,12 @@ struct CallbackClassifier {
       // placeholders added (ideally we'd remove them)
       // TODO: Remove known conditions and split the `if` statement
 
-      if (CallbackConditions.empty()) {
+      if (IsKnownConditionPath) {
+        // If we're on a known condition path, we can be lenient as we already
+        // know what block we're in and can therefore just add the conditional
+        // straight to it.
+        CurrentBlock->addNode(Statement);
+      } else if (CallbackConditions.empty()) {
         // Technically this has a similar problem, ie. the else could have
         // conditions that should be in either success/error
         CurrentBlock->addNode(Statement);
@@ -5311,19 +5322,45 @@ struct CallbackClassifier {
 
     if (ElseStmt) {
       if (auto *BS = dyn_cast<BraceStmt>(ElseStmt)) {
-        setNodes(ElseBlock, ThenBlock, NodesToPrint::inBraceStmt(BS));
+        // If this is a guard statement, we know that we'll always exit,
+        // allowing us to classify any additional nodes into the opposite block.
+        auto AlwaysExits = isa<GuardStmt>(Statement);
+        setNodes(ElseBlock, ThenBlock, NodesToPrint::inBraceStmt(BS),
+                 AlwaysExits);
       } else {
+        // If we reached here, we should have an else if statement. Given we
+        // know we're in the else of a known condition, temporarily flip the
+        // current block, and set that we know what path we're on.
+        llvm::SaveAndRestore<bool> CondScope(IsKnownConditionPath, true);
+        llvm::SaveAndRestore<ClassifiedBlock *> BlockScope(CurrentBlock,
+                                                           ElseBlock);
         classifyNodes(ArrayRef<ASTNode>(ElseStmt),
                       /*endCommentLoc*/ SourceLoc());
       }
     }
   }
 
+  /// Adds \p Nodes to \p Block, potentially flipping the current block if we
+  /// can determine that the nodes being added will cause control flow to leave
+  /// the scope.
+  ///
+  /// \param Block The block to add the nodes to.
+  /// \param OtherBlock The block for the opposing condition path.
+  /// \param Nodes The nodes to add.
+  /// \param AlwaysExitsScope Whether the nodes being added always exit the
+  /// scope, and therefore whether the current block should be flipped.
   void setNodes(ClassifiedBlock *Block, ClassifiedBlock *OtherBlock,
-                NodesToPrint Nodes) {
-    if (Nodes.hasTrailingReturnOrBreak()) {
-      CurrentBlock = OtherBlock;
+                NodesToPrint Nodes, bool AlwaysExitsScope = false) {
+    // Drop an explicit trailing 'return' or 'break' if we can.
+    bool HasTrailingReturnOrBreak = Nodes.hasTrailingReturnOrBreak();
+    if (HasTrailingReturnOrBreak)
       Nodes.dropTrailingReturnOrBreakIfPossible();
+
+    // If we know we're exiting the scope, we can set IsKnownConditionPath, as
+    // we know any future nodes should be classified into the other block.
+    if (HasTrailingReturnOrBreak || AlwaysExitsScope) {
+      CurrentBlock = OtherBlock;
+      IsKnownConditionPath = true;
       Block->addAllNodes(std::move(Nodes));
     } else {
       Block->addAllNodes(std::move(Nodes));

test/refactoring/ConvertAsync/convert_params_multi.swift

@@ -25,6 +25,13 @@ manyWithError { res1, res2, err in
 // MANYBOUND-NEXT: print("got error \(bad)")
 // MANYBOUND-NEXT: }
 
+// FIXME: This case is a little tricky: Being in the else block of 'if let str = res1'
+// should allow us to place 'if let i = res2' in the failure block. However, this
+// is a success condition, so we still place it in the success block. Really what
+// we need to do here is check to see if manyWithError has an existing async
+// alternative that still returns optional success values, and allow success
+// classification in that case. Otherwise, we'd probably be better off leaving
+// the condition unhandled, as it's not clear what the user is doing.
 // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=MANYUNBOUND-ERR %s
 manyWithError { res1, res2, err in
   print("before")

test/refactoring/ConvertAsync/convert_pattern.swift

@@ -121,9 +121,7 @@ func testPatterns() async throws {
   // FALLBACK-NEXT: guard let (str1, str2) = strs, str1 == "hi" else { fatalError() }
   // FALLBACK-NEXT: print(str1, str2, err)
 
-  // FIXME: Arguably we should be able to classify everything after the guard as
-  // a success path and avoid the fallback in this case.
-  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=FALLBACK2 %s
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=GUARD-AND-UNHANDLED %s
   stringTupleParam { strs, err in
     guard let (str1, str2) = strs else { fatalError() }
     print(str1, str2)
@@ -134,22 +132,17 @@ func testPatterns() async throws {
     }
   }
 
-  // FALLBACK2:      var strs: (String, String)? = nil
-  // FALLBACK2-NEXT: var err: Error? = nil
-  // FALLBACK2-NEXT: do {
-  // FALLBACK2-NEXT:   strs = try await stringTupleParam()
-  // FALLBACK2-NEXT: } catch {
-  // FALLBACK2-NEXT:   err = error
-  // FALLBACK2-NEXT: }
-  // FALLBACK2-EMPTY:
-  // FALLBACK2-NEXT: guard let (str1, str2) = strs else { fatalError() }
-  // FALLBACK2-NEXT: print(str1, str2)
-  // FALLBACK2-NEXT: if .random(), err == nil {
-  // FALLBACK2-NEXT:   print("yay")
-  // FALLBACK2-NEXT: } else {
-  // FALLBACK2-NEXT:   print("nay")
-  // FALLBACK2-NEXT: }
-
+  // GUARD-AND-UNHANDLED:      do {
+  // GUARD-AND-UNHANDLED-NEXT:   let (str1, str2) = try await stringTupleParam()
+  // GUARD-AND-UNHANDLED-NEXT:   print(str1, str2)
+  // GUARD-AND-UNHANDLED-NEXT:   if .random(), <#err#> == nil {
+  // GUARD-AND-UNHANDLED-NEXT:     print("yay")
+  // GUARD-AND-UNHANDLED-NEXT:   } else {
+  // GUARD-AND-UNHANDLED-NEXT:     print("nay")
+  // GUARD-AND-UNHANDLED-NEXT:   }
+  // GUARD-AND-UNHANDLED-NEXT: } catch let err {
+  // GUARD-AND-UNHANDLED-NEXT:   fatalError()
+  // GUARD-AND-UNHANDLED-NEXT: }
 
   // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=MIXED-BINDINGS %s
   stringTupleParam { strs, err in

test/refactoring/ConvertAsync/path_classification.swift

@@ -0,0 +1,234 @@
+// RUN: %empty-directory(%t)
+
+func simpleWithError(completion: (String?, Error?) -> Void) {}
+func simpleWithError() async throws -> String {}
+
+func testPathClassification() async throws {
+
+  // Both of these test cases should produce the same refactoring.
+
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION %s
+  simpleWithError { str, err in
+    if err == nil {
+      print("a")
+    } else if .random() {
+      print("b")
+    } else {
+      print("c")
+    }
+    if err != nil {
+      print("d")
+    } else if .random() {
+      print("e")
+    } else {
+      print("f")
+    }
+  }
+
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION %s
+  simpleWithError { str, err in
+    if let str = str {
+      print("a")
+    } else if .random() {
+      print("b")
+    } else {
+      print("c")
+    }
+    if str == nil {
+      print("d")
+    } else if .random() {
+      print("e")
+    } else {
+      print("f")
+    }
+  }
+
+  // ELSE-IF-CLASSIFICATION:      do {
+  // ELSE-IF-CLASSIFICATION-NEXT:   let str = try await simpleWithError()
+  // ELSE-IF-CLASSIFICATION-NEXT:   print("a")
+  // ELSE-IF-CLASSIFICATION-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION-NEXT:     print("e")
+  // ELSE-IF-CLASSIFICATION-NEXT:   } else {
+  // ELSE-IF-CLASSIFICATION-NEXT:     print("f")
+  // ELSE-IF-CLASSIFICATION-NEXT:   }
+  // ELSE-IF-CLASSIFICATION-NEXT: } catch let err {
+  // ELSE-IF-CLASSIFICATION-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION-NEXT:     print("b")
+  // ELSE-IF-CLASSIFICATION-NEXT:   } else {
+  // ELSE-IF-CLASSIFICATION-NEXT:     print("c")
+  // ELSE-IF-CLASSIFICATION-NEXT:   }
+  // ELSE-IF-CLASSIFICATION-NEXT:   print("d")
+  // ELSE-IF-CLASSIFICATION-NEXT: }
+
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION2 %s
+  simpleWithError { str, err in
+    if err == nil {
+      print("a")
+    } else if .random() {
+      print("b")
+    }
+    if .random() {
+      print("c")
+    }
+    if err != nil, .random() {
+      print("d")
+    }
+  }
+
+  // Make sure the classification of 'b' into the error block doesn't affect the
+  // handling of 'c'.
+  // ELSE-IF-CLASSIFICATION2:      do {
+  // ELSE-IF-CLASSIFICATION2-NEXT:   let str = try await simpleWithError()
+  // ELSE-IF-CLASSIFICATION2-NEXT:   print("a")
+  // ELSE-IF-CLASSIFICATION2-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION2-NEXT:     print("c")
+  // ELSE-IF-CLASSIFICATION2-NEXT:   }
+  // ELSE-IF-CLASSIFICATION2-NEXT: } catch let err {
+  // ELSE-IF-CLASSIFICATION2-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION2-NEXT:     print("b")
+  // ELSE-IF-CLASSIFICATION2-NEXT:   }
+  // ELSE-IF-CLASSIFICATION2-NEXT:   if <#err#> != nil, .random() {
+  // ELSE-IF-CLASSIFICATION2-NEXT:     print("d")
+  // ELSE-IF-CLASSIFICATION2-NEXT:   }
+  // ELSE-IF-CLASSIFICATION2-NEXT: }
+
+  // FIXME: This case is a little tricky: Being in the else block of 'err == nil'
+  // should allow us to place 'if let str = str' in the failure block. However, this
+  // is a success condition, so we still place it in the success block. Really what
+  // we need to do here is check to see if simpleWithError has an existing async
+  // alternative that still returns optional success values, and allow success
+  // classification in that case. Otherwise, we'd probably be better off leaving
+  // the condition unhandled, as it's not clear what the user is doing.
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION3 %s
+  simpleWithError { str, err in
+    if err == nil {
+      print("a")
+    } else if let str = str {
+      print("b")
+    } else {
+      print("c")
+    }
+  }
+
+  // ELSE-IF-CLASSIFICATION3:      do {
+  // ELSE-IF-CLASSIFICATION3-NEXT:   let str = try await simpleWithError()
+  // ELSE-IF-CLASSIFICATION3-NEXT:   print("a")
+  // ELSE-IF-CLASSIFICATION3-NEXT:   print("b")
+  // ELSE-IF-CLASSIFICATION3-NEXT: } catch let err {
+  // ELSE-IF-CLASSIFICATION3-NEXT:   print("c")
+  // ELSE-IF-CLASSIFICATION3-NEXT: }
+
+  // FIXME: Similar to the case above.
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION4 %s
+  simpleWithError { str, err in
+    if err == nil {
+      print("a")
+    } else if let str = str {
+      print("b")
+    } else if .random() {
+      print("c")
+    } else {
+      print("d")
+    }
+  }
+
+  // ELSE-IF-CLASSIFICATION4:      do {
+  // ELSE-IF-CLASSIFICATION4-NEXT:   let str = try await simpleWithError()
+  // ELSE-IF-CLASSIFICATION4-NEXT:   print("a")
+  // ELSE-IF-CLASSIFICATION4-NEXT:   print("b")
+  // ELSE-IF-CLASSIFICATION4-NEXT: } catch let err {
+  // ELSE-IF-CLASSIFICATION4-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION4-NEXT:     print("c")
+  // ELSE-IF-CLASSIFICATION4-NEXT:   } else {
+  // ELSE-IF-CLASSIFICATION4-NEXT:     print("d")
+  // ELSE-IF-CLASSIFICATION4-NEXT:   }
+  // ELSE-IF-CLASSIFICATION4-NEXT: }
+
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=ELSE-IF-CLASSIFICATION5 %s
+  simpleWithError { str, err in
+    if let err = err {
+      print("a")
+    } else if let str = str {
+      print("b")
+      return
+    } else if .random() {
+      print("c")
+    } else {
+      print("d")
+    }
+  }
+
+  // ELSE-IF-CLASSIFICATION5:      do {
+  // ELSE-IF-CLASSIFICATION5-NEXT:   let str = try await simpleWithError()
+  // ELSE-IF-CLASSIFICATION5-NEXT:   print("b")
+  // ELSE-IF-CLASSIFICATION5-NEXT: } catch let err {
+  // ELSE-IF-CLASSIFICATION5-NEXT:   print("a")
+  // ELSE-IF-CLASSIFICATION5-NEXT:   if .random() {
+  // ELSE-IF-CLASSIFICATION5-NEXT:     print("c")
+  // ELSE-IF-CLASSIFICATION5-NEXT:   } else {
+  // ELSE-IF-CLASSIFICATION5-NEXT:     print("d")
+  // ELSE-IF-CLASSIFICATION5-NEXT:   }
+  // ELSE-IF-CLASSIFICATION5-NEXT: }
+
+  // RN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=IF-LET-RETURN-CLASSIFICATION %s
+  simpleWithError { str, err in
+    if let str = str {
+      print("a")
+      return
+    }
+    if .random(), let err = err {
+      print("b")
+    } else {
+      print("c")
+    }
+  }
+
+  // IF-LET-RETURN-CLASSIFICATION-NEXT: do {
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:   let str = try await simpleWithError()
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:   print("a")
+  // IF-LET-RETURN-CLASSIFICATION-NEXT: } catch let err {
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:   if .random(), let err = <#err#> {
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:     print("b")
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:   } else {
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:     print("c")
+  // IF-LET-RETURN-CLASSIFICATION-NEXT:   }
+  // IF-LET-RETURN-CLASSIFICATION-NEXT: }
+
+  // RN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=GUARD-CLASSIFICATION %s
+  simpleWithError { str, err in
+    guard let str = str else {
+      print("a")
+      return
+    }
+    guard err == nil, .random() else {
+      print("b")
+      return
+    }
+    print("c")
+  }
+
+  // GUARD-CLASSIFICATION:      do {
+  // GUARD-CLASSIFICATION-NEXT:   let str = try await simpleWithError()
+  // GUARD-CLASSIFICATION-NEXT:   guard <#err#> == nil, .random() else {
+  // GUARD-CLASSIFICATION-NEXT:     print("b")
+  // GUARD-CLASSIFICATION-NEXT:     <#return#>
+  // GUARD-CLASSIFICATION-NEXT:   }
+  // GUARD-CLASSIFICATION-NEXT:   print("c")
+  // GUARD-CLASSIFICATION-NEXT: } catch let err {
+  // GUARD-CLASSIFICATION-NEXT:   print("a")
+  // GUARD-CLASSIFICATION-NEXT: }
+
+  // RUN: %refactor-check-compiles -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=SILLY-CLASSIFICATION %s
+  simpleWithError { str, err in
+    guard let str = str else { return }
+    guard let err = err else { return }
+    print("urr")
+  }
+
+  // In this case we just take whichever guard is last.
+  // SILLY-CLASSIFICATION:      do {
+  // SILLY-CLASSIFICATION-NEXT:   let str = try await simpleWithError()
+  // SILLY-CLASSIFICATION-NEXT: } catch let err {
+  // SILLY-CLASSIFICATION-NEXT:   print("urr")
+  // SILLY-CLASSIFICATION-NEXT: }
+}