Merge pull request #38052 from ahoppen/pr/use-closure-label-for-async-return-type

[Refactoring] Use internal completion handler labels for async function's return type

Author: ahoppen

lib/IDE/Refactoring.cpp

@@ -4079,6 +4079,16 @@ class HandlerResult {
 /// single parameter of `Result` type.
 enum class HandlerType { INVALID, PARAMS, RESULT };
 
+/// A single return type of a refactored async function. If the async function
+/// returns a tuple, each element of the tuple (represented by a \c
+/// LabeledReturnType) might have a label, otherwise the \p Label is empty.
+struct LabeledReturnType {
+  Identifier Label;
+  swift::Type Ty;
+
+  LabeledReturnType(Identifier Label, swift::Type Ty) : Label(Label), Ty(Ty) {}
+};
+
 /// Given a function with an async alternative (or one that *could* have an
 /// async alternative), stores information about the completion handler.
 /// The completion handler can be either a variable (which includes a parameter)
@@ -4327,12 +4337,26 @@ struct AsyncHandlerDesc {
     }
   }
 
+  /// If the async function returns a tuple, the label of the \p Index -th
+  /// element in the returned tuple. If the function doesn't return a tuple or
+  /// the element is unlabeled, an empty identifier is returned.
+  Identifier getAsyncReturnTypeLabel(size_t Index) const {
+    assert(Index < getSuccessParams().size());
+    if (getSuccessParams().size() <= 1) {
+      // There can't be any labels if the async function doesn't return a tuple.
+      return Identifier();
+    } else {
+      return getSuccessParams()[Index].getInternalLabel();
+    }
+  }
+
   /// Gets the return value types for the async equivalent of this handler.
-  ArrayRef<swift::Type>
-  getAsyncReturnTypes(SmallVectorImpl<swift::Type> &Scratch) const {
-    for (auto &Param : getSuccessParams()) {
-      auto Ty = Param.getParameterType();
-      Scratch.push_back(getSuccessParamAsyncReturnType(Ty));
+  ArrayRef<LabeledReturnType>
+  getAsyncReturnTypes(SmallVectorImpl<LabeledReturnType> &Scratch) const {
+    for (size_t I = 0; I < getSuccessParams().size(); ++I) {
+      auto Ty = getSuccessParams()[I].getParameterType();
+      Scratch.emplace_back(getAsyncReturnTypeLabel(I),
+                           getSuccessParamAsyncReturnType(Ty));
     }
     return Scratch;
   }
@@ -6371,7 +6395,7 @@ class AsyncConverter : private SourceEntityWalker {
       return;
     }
 
-    SmallVector<Type, 2> Scratch;
+    SmallVector<LabeledReturnType, 2> Scratch;
     auto ReturnTypes = TopHandler.getAsyncReturnTypes(Scratch);
     if (ReturnTypes.empty()) {
       OS << " ";
@@ -6385,7 +6409,14 @@ class AsyncConverter : private SourceEntityWalker {
         OS << "(";
 
       llvm::interleave(
-          ReturnTypes, [&](Type Ty) { Ty->print(OS); }, [&]() { OS << ", "; });
+          ReturnTypes,
+          [&](LabeledReturnType TypeAndLabel) {
+            if (!TypeAndLabel.Label.empty()) {
+              OS << TypeAndLabel.Label << tok::colon << " ";
+            }
+            TypeAndLabel.Ty->print(OS);
+          },
+          [&]() { OS << ", "; });
 
       if (ReturnTypes.size() > 1)
         OS << ")";
@@ -7164,7 +7195,14 @@ class AsyncConverter : private SourceEntityWalker {
         //     completion(result.0, result.1)
         //   }
         // }
-        OS << ResultName << tok::period << Index;
+        OS << ResultName << tok::period;
+
+        auto Label = HandlerDesc.getAsyncReturnTypeLabel(Index);
+        if (!Label.empty()) {
+          OS << Label;
+        } else {
+          OS << Index;
+        }
       } else {
         OS << ResultName;
       }
@@ -7212,9 +7250,14 @@ class AsyncConverter : private SourceEntityWalker {
   /// returned results via a completion handler described by \p HandlerDesc.
   void addAsyncFuncReturnType(const AsyncHandlerDesc &HandlerDesc) {
     // Type or (Type1, Type2, ...)
-    SmallVector<Type, 2> Scratch;
+    SmallVector<LabeledReturnType, 2> Scratch;
     addTupleOf(HandlerDesc.getAsyncReturnTypes(Scratch), OS,
-               [&](auto Ty) { Ty->print(OS); });
+               [&](LabeledReturnType LabelAndType) {
+                 if (!LabelAndType.Label.empty()) {
+                   OS << LabelAndType.Label << tok::colon << " ";
+                 }
+                 LabelAndType.Ty->print(OS);
+               });
   }
 
   /// If \p FD is generic, adds a type annotation with the return type of the

test/refactoring/ConvertAsync/labeled_closure_params.swift

@@ -0,0 +1,70 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %refactor-check-compiles -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=MULTIPLE-LABELED-RESULTS %s
+func mutlipleLabeledResults(completion: (_ first: String, _ second: String) -> Void) { }
+// MULTIPLE-LABELED-RESULTS: {
+// MULTIPLE-LABELED-RESULTS-NEXT: async {
+// MULTIPLE-LABELED-RESULTS-NEXT: let result = await mutlipleLabeledResults()
+// MULTIPLE-LABELED-RESULTS-NEXT: completion(result.first, result.second)
+// MULTIPLE-LABELED-RESULTS-NEXT: }
+// MULTIPLE-LABELED-RESULTS-NEXT: }
+// MULTIPLE-LABELED-RESULTS: func mutlipleLabeledResults() async -> (first: String, second: String) { }
+
+// RUN: %refactor-check-compiles -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=MIXED-LABELED-RESULTS %s
+func mixedLabeledResult(completion: (_ first: String, String) -> Void) { }
+// MIXED-LABELED-RESULTS: {
+// MIXED-LABELED-RESULTS-NEXT: async {
+// MIXED-LABELED-RESULTS-NEXT: let result = await mixedLabeledResult()
+// MIXED-LABELED-RESULTS-NEXT: completion(result.first, result.1)
+// MIXED-LABELED-RESULTS-NEXT: }
+// MIXED-LABELED-RESULTS-NEXT: }
+// MIXED-LABELED-RESULTS: func mixedLabeledResult() async -> (first: String, String) { }
+
+// RUN: %refactor-check-compiles -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=SINGLE-LABELED-RESULT %s
+func singleLabeledResult(completion: (_ first: String) -> Void) { }
+// SINGLE-LABELED-RESULT: {
+// SINGLE-LABELED-RESULT-NEXT: async {
+// SINGLE-LABELED-RESULT-NEXT: let result = await singleLabeledResult()
+// SINGLE-LABELED-RESULT-NEXT: completion(result)
+// SINGLE-LABELED-RESULT-NEXT: }
+// SINGLE-LABELED-RESULT-NEXT: }
+// SINGLE-LABELED-RESULT: func singleLabeledResult() async -> String { }
+
+// RUN: %refactor-check-compiles -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=SINGLE-LABELED-RESULT-WITH-ERROR %s
+func singleLabeledResultWithError(completion: (_ first: String?, _ error: Error?) -> Void) { }
+// SINGLE-LABELED-RESULT-WITH-ERROR: {
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: async {
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: do {
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: let result = try await singleLabeledResultWithError()
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: completion(result, nil)
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: } catch {
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: completion(nil, error)
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// SINGLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// SINGLE-LABELED-RESULT-WITH-ERROR: func singleLabeledResultWithError() async throws -> String { }
+
+// RUN: %refactor-check-compiles -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=MULTIPLE-LABELED-RESULT-WITH-ERROR %s
+func multipleLabeledResultWithError(completion: (_ first: String?, _ second: String?, _ error: Error?) -> Void) { }
+// MULTIPLE-LABELED-RESULT-WITH-ERROR: {
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: async {
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: do {
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: let result = try await multipleLabeledResultWithError()
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: completion(result.first, result.second, nil)
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: } catch {
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: completion(nil, nil, error)
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// MULTIPLE-LABELED-RESULT-WITH-ERROR-NEXT: }
+// MULTIPLE-LABELED-RESULT-WITH-ERROR: func multipleLabeledResultWithError() async throws -> (first: String, second: String) { }
+
+func testConvertCall() {
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=CONVERT-CALL %s
+  mutlipleLabeledResults() { (a, b) in
+    print(a)
+    print(b)
+  }
+  // CONVERT-CALL: let (a, b) = await mutlipleLabeledResults()
+  // CONVERT-CALL-NEXT: print(a)
+  // CONVERT-CALL-NEXT: print(b)
+}