[Refactoring] Improve Void handling for async conversion

hamishknight · hamishknight · commit c1e4dc3d82cf · 2021-04-30T01:09:04.000+01:00
When converting a function with a completion handler
that has a Void success parameter, e.g
`(Void?, Error?) -&gt; Void`, or more likely a
`Result&lt;Void, Error&gt;` parameter, make sure to omit
the `-&gt; Void` from the resulting async function
conversion.

In addition, strip any Void bindings from an async
function call, and any explicit Void return values
from inside the async function.

Resolves rdar://75189289
diff --git a/lib/IDE/Refactoring.cpp b/lib/IDE/Refactoring.cpp
@@ -4141,6 +4141,14 @@ struct AsyncHandlerDesc {
     return Ty->getParams();
   }
 
+  /// Retrieve the parameters relevant to a successful return from the
+  /// completion handler. This drops the Error parameter if present.
+  ArrayRef<AnyFunctionType::Param> getSuccessParams() const {
+    if (HasError && Type == HandlerType::PARAMS)
+      return params().drop_back();
+    return params();
+  }
+
   /// The `CallExpr` if the given node is a call to the `Handler`
   CallExpr *getAsHandlerCall(ASTNode Node) const {
     if (!isValid())
@@ -4162,13 +4170,16 @@ struct AsyncHandlerDesc {
     auto Args = ArgList.ref();
 
     if (Type == HandlerType::PARAMS) {
-      if (!HasError)
-        return HandlerResult(Args);
-
-      if (!isa<NilLiteralExpr>(Args.back()))
+      // If there's an error parameter and the user isn't passing nil to it,
+      // assume this is the error path.
+      if (HasError && !isa<NilLiteralExpr>(Args.back()))
         return HandlerResult(Args.back(), true);
 
-      return HandlerResult(Args.drop_back());
+      // We can drop the args altogether if they're just Void.
+      if (willAsyncReturnVoid())
+        return HandlerResult();
+
+      return HandlerResult(HasError ? Args.drop_back() : Args);
     } else if (Type == HandlerType::RESULT) {
       if (Args.size() != 1)
         return HandlerResult(Args);
@@ -4187,12 +4198,62 @@ struct AsyncHandlerDesc {
         return HandlerResult(Args);
 
       auto ResultArgList = callArgs(ResultCE);
-      return HandlerResult(ResultArgList.ref()[0],
-                           D->getNameStr() == StringRef("failure"));
+      auto isFailure = D->getNameStr() == StringRef("failure");
+
+      // We can drop the arg altogether if it's just Void.
+      if (!isFailure && willAsyncReturnVoid())
+        return HandlerResult();
+
+      // Otherwise the arg gets the .success() or .failure() call dropped.
+      return HandlerResult(ResultArgList.ref()[0], isFailure);
     }
 
     llvm_unreachable("Unhandled result type");
   }
+
+  // Convert the type of a success parameter in the completion handler function
+  // to a return type suitable for an async function. If there is an error
+  // parameter present e.g (T?, Error?) -> Void, this unwraps a level of
+  // optionality from T?. If this is a Result<T, U> type, returns the success
+  // type T.
+  swift::Type getSuccessParamAsyncReturnType(swift::Type Ty) const {
+    switch (Type) {
+    case HandlerType::PARAMS: {
+      // If there's an Error parameter in the handler, the success branch can
+      // be unwrapped.
+      if (HasError)
+        Ty = Ty->lookThroughSingleOptionalType();
+
+      return Ty;
+    }
+    case HandlerType::RESULT: {
+      // Result<T, U> maps to T.
+      return Ty->castTo<BoundGenericType>()->getGenericArgs()[0];
+    }
+    case HandlerType::INVALID:
+      llvm_unreachable("Invalid handler type");
+    }
+  }
+
+  /// 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));
+    }
+    return Scratch;
+  }
+
+  /// Whether the async equivalent of this handler returns Void.
+  bool willAsyncReturnVoid() const {
+    // If all of the success params will be converted to Void return types,
+    // this will be a Void async function.
+    return llvm::all_of(getSuccessParams(), [&](auto &param) {
+      auto Ty = param.getParameterType();
+      return getSuccessParamAsyncReturnType(Ty)->isVoid();
+    });
+  }
 };
 
 enum class ConditionType { INVALID, NIL, NOT_NIL };
@@ -4911,42 +4972,24 @@ class AsyncConverter : private SourceEntityWalker {
       return;
     }
 
-    auto HandlerParams = TopHandler.params();
-    if (TopHandler.Type == HandlerType::PARAMS && TopHandler.HasError) {
-      HandlerParams = HandlerParams.drop_back();
-    }
-
-    if (HandlerParams.empty()) {
+    SmallVector<Type, 2> Scratch;
+    auto ReturnTypes = TopHandler.getAsyncReturnTypes(Scratch);
+    if (ReturnTypes.empty()) {
       OS << " ";
       return;
     }
 
-    OS << " -> ";
+    // Print the function result type, making sure to omit a '-> Void' return.
+    if (!TopHandler.willAsyncReturnVoid()) {
+      OS << " -> ";
+      if (ReturnTypes.size() > 1)
+        OS << "(";
 
-    if (HandlerParams.size() > 1) {
-      OS << "(";
-    }
-    for (size_t I = 0, E = HandlerParams.size(); I < E; ++I) {
-      if (I > 0) {
-        OS << ", ";
-      }
+      llvm::interleave(
+          ReturnTypes, [&](Type Ty) { Ty->print(OS); }, [&]() { OS << ", "; });
 
-      auto &Param = HandlerParams[I];
-      if (TopHandler.Type == HandlerType::PARAMS) {
-        Type ToPrint = Param.getPlainType();
-        if (TopHandler.HasError)
-          ToPrint = ToPrint->lookThroughSingleOptionalType();
-        ToPrint->print(OS);
-      } else if (TopHandler.Type == HandlerType::RESULT) {
-        auto ResultTy = Param.getPlainType()->getAs<BoundGenericType>();
-        assert(ResultTy && "Result must have generic type");
-        ResultTy->getGenericArgs()[0]->print(OS);
-      } else {
-        llvm_unreachable("Unhandled handler type");
-      }
-    }
-    if (HandlerParams.size() > 1) {
-      OS << ")";
+      if (ReturnTypes.size() > 1)
+        OS << ")";
     }
 
     if (FD->hasBody())
@@ -5073,8 +5116,7 @@ class AsyncConverter : private SourceEntityWalker {
         addFallbackVars(CallbackParams->getArray(), Blocks);
         addDo();
         addAwaitCall(CE, ArgList.ref(), Blocks.SuccessBlock, SuccessParams,
-                     /*HasError=*/HandlerDesc.HasError,
-                     /*AddDeclarations=*/!HandlerDesc.HasError);
+                     HandlerDesc, /*AddDeclarations=*/!HandlerDesc.HasError);
         addFallbackCatch(ErrParam);
         OS << "\n";
         convertNodes(CallbackBody);
@@ -5107,10 +5149,9 @@ class AsyncConverter : private SourceEntityWalker {
 
     setNames(Blocks.SuccessBlock, SuccessParams);
     addAwaitCall(CE, ArgList.ref(), Blocks.SuccessBlock, SuccessParams,
-                 /*HasError=*/HandlerDesc.HasError,
-                 /*AddDeclarations=*/true);
+                 HandlerDesc, /*AddDeclarations=*/true);
 
-    prepareNamesForBody(HandlerDesc.Type, SuccessParams, ErrParams);
+    prepareNamesForBody(HandlerDesc, SuccessParams, ErrParams);
     convertNodes(Blocks.SuccessBlock.nodes());
 
     if (RequireDo) {
@@ -5120,7 +5161,7 @@ class AsyncConverter : private SourceEntityWalker {
                HandlerDesc.Type != HandlerType::RESULT);
       addCatch(ErrParam);
 
-      prepareNamesForBody(HandlerDesc.Type, ErrParams, SuccessParams);
+      prepareNamesForBody(HandlerDesc, ErrParams, SuccessParams);
       addCatchBody(ErrParam, Blocks.ErrorBlock);
     }
 
@@ -5129,9 +5170,11 @@ class AsyncConverter : private SourceEntityWalker {
 
   void addAwaitCall(const CallExpr *CE, ArrayRef<Expr *> Args,
                     const ClassifiedBlock &SuccessBlock,
-                    ArrayRef<const ParamDecl *> SuccessParams, bool HasError,
-                    bool AddDeclarations) {
-    if (!SuccessParams.empty()) {
+                    ArrayRef<const ParamDecl *> SuccessParams,
+                    const AsyncHandlerDesc &HandlerDesc, bool AddDeclarations) {
+    // Print the bindings to match the completion handler success parameters,
+    // making sure to omit in the case of a Void return.
+    if (!SuccessParams.empty() && !HandlerDesc.willAsyncReturnVoid()) {
       if (AddDeclarations) {
         if (SuccessBlock.allLet()) {
           OS << tok::kw_let;
@@ -5153,7 +5196,7 @@ class AsyncConverter : private SourceEntityWalker {
       OS << " " << tok::equal << " ";
     }
 
-    if (HasError) {
+    if (HandlerDesc.HasError) {
       OS << tok::kw_try << " ";
     }
     OS << "await ";
@@ -5199,16 +5242,21 @@ class AsyncConverter : private SourceEntityWalker {
     OS << "\n" << tok::r_brace;
   }
 
-  void prepareNamesForBody(HandlerType ResultType,
+  void prepareNamesForBody(const AsyncHandlerDesc &HandlerDesc,
                            ArrayRef<const ParamDecl *> CurrentParams,
                            ArrayRef<const ParamDecl *> OtherParams) {
-    switch (ResultType) {
+    switch (HandlerDesc.Type) {
     case HandlerType::PARAMS:
       for (auto *Param : CurrentParams) {
-        if (Param->getType()->getOptionalObjectType()) {
+        auto Ty = Param->getType();
+        if (Ty->getOptionalObjectType()) {
           Unwraps.insert(Param);
           Placeholders.insert(Param);
         }
+        // Void parameters get omitted where possible, so turn any reference
+        // into a placeholder, as its usage is unlikely what the user wants.
+        if (HandlerDesc.getSuccessParamAsyncReturnType(Ty)->isVoid())
+          Placeholders.insert(Param);
       }
       // Use of the other params is invalid within the current body
       Placeholders.insert(OtherParams.begin(), OtherParams.end());
diff --git a/test/refactoring/ConvertAsync/basic.swift b/test/refactoring/ConvertAsync/basic.swift
@@ -190,6 +190,18 @@ func blockConvention(completion: @convention(block) () -> Void) { }
 func cConvention(completion: @convention(c) () -> Void) { }
 // C-CONVENTION: func cConvention() async { }
 
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-HANDLER %s
+func voidCompletion(completion: (Void) -> Void) {}
+// VOID-HANDLER: func voidCompletion() async {}
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-AND-ERROR-HANDLER %s
+func voidAndErrorCompletion(completion: (Void?, Error?) -> Void) {}
+// VOID-AND-ERROR-HANDLER: func voidAndErrorCompletion() async throws {}
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix TOO-MUCH-VOID-AND-ERROR-HANDLER %s
+func tooMuchVoidAndErrorCompletion(completion: (Void?, Void?, Error?) -> Void) {}
+// TOO-MUCH-VOID-AND-ERROR-HANDLER: func tooMuchVoidAndErrorCompletion() async throws {}
+
 // 2. Check that the various ways to call a function (and the positions the
 //    refactoring is called from) are handled correctly
 
@@ -348,5 +360,26 @@ func testCalls() {
   }
   // C-CONVENTION-CALL: await cConvention(){{$}}
   // C-CONVENTION-CALL-NEXT: {{^}}print("cConvention")
+
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=VOID-AND-ERROR-CALL %s
+  voidAndErrorCompletion { v, err in
+    print("void and error completion \(v)")
+  }
+  // VOID-AND-ERROR-CALL: {{^}}try await voidAndErrorCompletion(){{$}}
+  // VOID-AND-ERROR-CALL: {{^}}print("void and error completion \(<#v#>)"){{$}}
+
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=VOID-AND-ERROR-CALL2 %s
+  voidAndErrorCompletion { _, err in
+    print("void and error completion 2")
+  }
+  // VOID-AND-ERROR-CALL2: {{^}}try await voidAndErrorCompletion(){{$}}
+  // VOID-AND-ERROR-CALL2: {{^}}print("void and error completion 2"){{$}}
+
+  // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):3 | %FileCheck -check-prefix=VOID-AND-ERROR-CALL3 %s
+  tooMuchVoidAndErrorCompletion { v, v1, err in
+    print("void and error completion 3")
+  }
+  // VOID-AND-ERROR-CALL3: {{^}}try await tooMuchVoidAndErrorCompletion(){{$}}
+  // VOID-AND-ERROR-CALL3: {{^}}print("void and error completion 3"){{$}}
 }
 // CONVERT-FUNC: {{^}}}
diff --git a/test/refactoring/ConvertAsync/convert_function.swift b/test/refactoring/ConvertAsync/convert_function.swift
@@ -148,3 +148,51 @@ func asyncUnhandledCompletion(_ completion: (String) -> Void) {
 // ASYNC-UNHANDLED-NEXT: return "bad"
 // ASYNC-UNHANDLED-NEXT: }
 // ASYNC-UNHANDLED-NEXT: }
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-AND-ERROR-HANDLER %s
+func voidAndErrorCompletion(completion: (Void?, Error?) -> Void) {
+  if .random() {
+    completion((), nil) // Make sure we drop the ()
+  } else {
+    completion(nil, CustomError.Bad)
+  }
+}
+// VOID-AND-ERROR-HANDLER:      func voidAndErrorCompletion() async throws {
+// VOID-AND-ERROR-HANDLER-NEXT:   if .random() {
+// VOID-AND-ERROR-HANDLER-NEXT:     return // Make sure we drop the ()
+// VOID-AND-ERROR-HANDLER-NEXT:   } else {
+// VOID-AND-ERROR-HANDLER-NEXT:     throw CustomError.Bad
+// VOID-AND-ERROR-HANDLER-NEXT:   }
+// VOID-AND-ERROR-HANDLER-NEXT: }
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix TOO-MUCH-VOID-AND-ERROR-HANDLER %s
+func tooMuchVoidAndErrorCompletion(completion: (Void?, Void?, Error?) -> Void) {
+  if .random() {
+    completion((), (), nil) // Make sure we drop the ()s
+  } else {
+    completion(nil, nil, CustomError.Bad)
+  }
+}
+// TOO-MUCH-VOID-AND-ERROR-HANDLER:      func tooMuchVoidAndErrorCompletion() async throws {
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT:   if .random() {
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT:     return // Make sure we drop the ()s
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT:   } else {
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT:     throw CustomError.Bad
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT:   }
+// TOO-MUCH-VOID-AND-ERROR-HANDLER-NEXT: }
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-RESULT-HANDLER %s
+func voidResultCompletion(completion: (Result<Void, Error>) -> Void) {
+  if .random() {
+    completion(.success(())) // Make sure we drop the .success(())
+  } else {
+    completion(.failure(CustomError.Bad))
+  }
+}
+// VOID-RESULT-HANDLER:      func voidResultCompletion() async throws {
+// VOID-RESULT-HANDLER-NEXT:   if .random() {
+// VOID-RESULT-HANDLER-NEXT:     return // Make sure we drop the .success(())
+// VOID-RESULT-HANDLER-NEXT:   } else {
+// VOID-RESULT-HANDLER-NEXT:     throw CustomError.Bad
+// VOID-RESULT-HANDLER-NEXT:   }
+// VOID-RESULT-HANDLER-NEXT: }
diff --git a/test/refactoring/ConvertAsync/convert_result.swift b/test/refactoring/ConvertAsync/convert_result.swift
@@ -2,6 +2,14 @@ func simple(_ completion: (Result<String, Error>) -> Void) { }
 func noError(_ completion: (Result<String, Never>) -> Void) { }
 func test(_ str: String) -> Bool { return false }
 
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-RESULT %s
+func voidResult(completion: (Result<Void, Never>) -> Void) {}
+// VOID-RESULT: func voidResult() async {}
+
+// RUN: %refactor -add-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix VOID-AND-ERROR-RESULT %s
+func voidAndErrorResult(completion: (Result<Void, Error>) -> Void) {}
+// VOID-AND-ERROR-RESULT: func voidAndErrorResult() async throws {}
+
 // RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=SIMPLE %s
 simple { res in
   print("result \(res)")
@@ -301,3 +309,16 @@ simple { res in
 // NESTEDBREAK-NEXT: print("after")
 // NESTEDBREAK-NOT: }
 
+// RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=VOID-RESULT-CALL %s
+voidResult { res in
+  print(res)
+}
+// VOID-RESULT-CALL: {{^}}await voidResult()
+// VOID-RESULT-CALL: {{^}}print(<#res#>)
+
+// RUN: %refactor -convert-call-to-async-alternative -dump-text -source-filename %s -pos=%(line+1):1 | %FileCheck -check-prefix=VOID-AND-ERROR-RESULT-CALL %s
+voidAndErrorResult { res in
+  print(res)
+}
+// VOID-AND-ERROR-RESULT-CALL: {{^}}try await voidAndErrorResult()
+// VOID-AND-ERROR-RESULT-CALL: {{^}}print(<#res#>)
