Merge pull request swiftlang#37491 from hamishknight/so-in-sync

Author: hamishknight

include/swift/AST/ParameterList.h

@@ -76,7 +76,10 @@ class alignas(ParamDecl *) ParameterList final :
   iterator end() { return getArray().end(); }
   const_iterator begin() const { return getArray().begin(); }
   const_iterator end() const { return getArray().end(); }
-  
+
+  ParamDecl *front() const { return getArray().front(); }
+  ParamDecl *back() const { return getArray().back(); }
+
   MutableArrayRef<ParamDecl*> getArray() {
     return {getTrailingObjects<ParamDecl*>(), numParameters};
   }

include/swift/IDE/RefactoringKinds.def

@@ -60,6 +60,8 @@ CURSOR_REFACTORING(ConvertToAsync, "Convert Function to Async", convert.func-to-
 
 CURSOR_REFACTORING(AddAsyncAlternative, "Add Async Alternative", add.async-alternative)
 
+CURSOR_REFACTORING(AddAsyncWrapper, "Add Async Wrapper", add.async-wrapper)
+
 RANGE_REFACTORING(ExtractExpr, "Extract Expression", extract.expr)
 
 RANGE_REFACTORING(ExtractFunction, "Extract Method", extract.function)

lib/IDE/Refactoring.cpp

@@ -4195,6 +4195,13 @@ struct AsyncHandlerDesc {
     return params();
   }
 
+  /// If the completion handler has an Error parameter, return it.
+  Optional<AnyFunctionType::Param> getErrorParam() const {
+    if (HasError && Type == HandlerType::PARAMS)
+      return params().back();
+    return None;
+  }
+
   /// Get the type of the error that will be thrown by the \c async method or \c
   /// None if the completion handler doesn't accept an error parameter.
   /// This may be more specialized than the generic 'Error' type if the
@@ -5397,6 +5404,41 @@ class AsyncConverter : private SourceEntityWalker {
     return true;
   }
 
+  /// Creates an async alternative function that forwards onto the completion
+  /// handler function through
+  /// withCheckedContinuation/withCheckedThrowingContinuation.
+  bool createAsyncWrapper() {
+    assert(Buffer.empty() && "AsyncConverter can only be used once");
+    auto *FD = cast<FuncDecl>(StartNode.get<Decl *>());
+
+    // First add the new async function declaration.
+    addFuncDecl(FD);
+    OS << tok::l_brace << "\n";
+
+    // Then add the body.
+    OS << tok::kw_return << " ";
+    if (TopHandler.HasError)
+      OS << tok::kw_try << " ";
+
+    OS << "await ";
+
+    // withChecked[Throwing]Continuation { cont in
+    if (TopHandler.HasError) {
+      OS << "withCheckedThrowingContinuation";
+    } else {
+      OS << "withCheckedContinuation";
+    }
+    OS << " " << tok::l_brace << " cont " << tok::kw_in << "\n";
+
+    // fnWithHandler(args...) { ... }
+    auto ClosureStr = getAsyncWrapperCompletionClosure("cont", TopHandler);
+    addForwardingCallTo(FD, TopHandler, /*HandlerReplacement*/ ClosureStr);
+
+    OS << tok::r_brace << "\n"; // end continuation closure
+    OS << tok::r_brace << "\n"; // end function body
+    return true;
+  }
+
   void replace(ASTNode Node, SourceEditConsumer &EditConsumer,
                SourceLoc StartOverride = SourceLoc()) {
     SourceRange Range = Node.getSourceRange();
@@ -5432,6 +5474,130 @@ class AsyncConverter : private SourceEntityWalker {
     return TopHandler.isValid();
   }
 
+  /// Prints a tuple of elements, or a lone single element if only one is
+  /// present, using the provided printing function.
+  template <typename T, typename PrintFn>
+  void addTupleOf(ArrayRef<T> Elements, llvm::raw_ostream &OS,
+                  PrintFn PrintElt) {
+    if (Elements.size() == 1) {
+      PrintElt(Elements[0]);
+      return;
+    }
+    OS << tok::l_paren;
+    llvm::interleave(Elements, PrintElt, [&]() { OS << tok::comma << " "; });
+    OS << tok::r_paren;
+  }
+
+  /// Retrieve the completion handler closure argument for an async wrapper
+  /// function.
+  std::string
+  getAsyncWrapperCompletionClosure(StringRef ContName,
+                                   const AsyncHandlerParamDesc &HandlerDesc) {
+    std::string OutputStr;
+    llvm::raw_string_ostream OS(OutputStr);
+
+    OS << " " << tok::l_brace; // start closure
+
+    // Prepare parameter names for the closure.
+    auto SuccessParams = HandlerDesc.getSuccessParams();
+    SmallVector<SmallString<4>, 2> SuccessParamNames;
+    for (auto idx : indices(SuccessParams)) {
+      SuccessParamNames.emplace_back("res");
+
+      // If we have multiple success params, number them e.g res1, res2...
+      if (SuccessParams.size() > 1)
+        SuccessParamNames.back().append(std::to_string(idx + 1));
+    }
+    Optional<SmallString<4>> ErrName;
+    if (HandlerDesc.getErrorParam())
+      ErrName.emplace("err");
+
+    auto HasAnyParams = !SuccessParamNames.empty() || ErrName;
+    if (HasAnyParams)
+      OS << " ";
+
+    // res1, res2
+    llvm::interleave(
+        SuccessParamNames, [&](auto Name) { OS << Name; },
+        [&]() { OS << tok::comma << " "; });
+
+    // , err
+    if (ErrName) {
+      if (!SuccessParamNames.empty())
+        OS << tok::comma << " ";
+
+      OS << *ErrName;
+    }
+    if (HasAnyParams)
+      OS << " " << tok::kw_in;
+
+    OS << "\n";
+
+    // The closure body.
+    switch (HandlerDesc.Type) {
+    case HandlerType::PARAMS: {
+      // For a (Success?, Error?) -> Void handler, we do an if let on the error.
+      if (ErrName) {
+        // if let err = err {
+        OS << tok::kw_if << " " << tok::kw_let << " ";
+        OS << *ErrName << " " << tok::equal << " " << *ErrName << " ";
+        OS << tok::l_brace << "\n";
+
+        // cont.resume(throwing: err)
+        OS << ContName << tok::period << "resume" << tok::l_paren;
+        OS << "throwing" << tok::colon << " " << *ErrName;
+        OS << tok::r_paren << "\n";
+
+        // return }
+        OS << tok::kw_return << "\n";
+        OS << tok::r_brace << "\n";
+      }
+
+      // If we have any success params that we need to unwrap, insert a guard.
+      for (auto Idx : indices(SuccessParamNames)) {
+        auto &Name = SuccessParamNames[Idx];
+        auto ParamTy = SuccessParams[Idx].getParameterType();
+        if (!HandlerDesc.shouldUnwrap(ParamTy))
+          continue;
+
+        // guard let res = res else {
+        OS << tok::kw_guard << " " << tok::kw_let << " ";
+        OS << Name << " " << tok::equal << " " << Name << " " << tok::kw_else;
+        OS << " " << tok::l_brace << "\n";
+
+        // fatalError(...)
+        OS << "fatalError" << tok::l_paren;
+        OS << "\"Expected non-nil success param '" << Name;
+        OS << "' for nil error\"";
+        OS << tok::r_paren << "\n";
+
+        // End guard.
+        OS << tok::r_brace << "\n";
+      }
+
+      // cont.resume(returning: (res1, res2, ...))
+      OS << ContName << tok::period << "resume" << tok::l_paren;
+      OS << "returning" << tok::colon << " ";
+      addTupleOf(llvm::makeArrayRef(SuccessParamNames), OS,
+                 [&](auto Ref) { OS << Ref; });
+      OS << tok::r_paren << "\n";
+      break;
+    }
+    case HandlerType::RESULT: {
+      // cont.resume(with: res)
+      assert(SuccessParamNames.size() == 1);
+      OS << ContName << tok::period << "resume" << tok::l_paren;
+      OS << "with" << tok::colon << " " << SuccessParamNames[0];
+      OS << tok::r_paren << "\n";
+      break;
+    }
+    case HandlerType::INVALID:
+      llvm_unreachable("Should not have an invalid handler here");
+    }
+
+    OS << tok::r_brace << "\n"; // end closure
+    return OutputStr;
+  }
 
   /// Retrieves the location for the start of a comment attached to the token
   /// at the provided location, or the location itself if there is no comment.
@@ -6075,16 +6241,9 @@ class AsyncConverter : private SourceEntityWalker {
         }
         OS << " ";
       }
-      if (SuccessParams.size() > 1)
-        OS << tok::l_paren;
-      OS << newNameFor(SuccessParams.front());
-      for (const auto Param : SuccessParams.drop_front()) {
-        OS << tok::comma << " ";
-        OS << newNameFor(Param);
-      }
-      if (SuccessParams.size() > 1) {
-        OS << tok::r_paren;
-      }
+      // 'res =' or '(res1, res2, ...) ='
+      addTupleOf(SuccessParams, OS,
+                 [&](auto &Param) { OS << newNameFor(Param); });
       OS << " " << tok::equal << " ";
     }
 
@@ -6271,22 +6430,46 @@ class AsyncConverter : private SourceEntityWalker {
   /// 'await' keyword.
   void addCallToAsyncMethod(const FuncDecl *FD,
                             const AsyncHandlerDesc &HandlerDesc) {
+    // The call to the async function is the same as the call to the old
+    // completion handler function, minus the completion handler arg.
+    addForwardingCallTo(FD, HandlerDesc, /*HandlerReplacement*/ "");
+  }
+
+  /// Adds a forwarding call to the old completion handler function, with
+  /// \p HandlerReplacement that allows for a custom replacement or, if empty,
+  /// removal of the completion handler closure.
+  void addForwardingCallTo(
+      const FuncDecl *FD, const AsyncHandlerDesc &HandlerDesc,
+      StringRef HandlerReplacement, bool CanUseTrailingClosure = true) {
     OS << FD->getBaseName() << tok::l_paren;
-    bool FirstParam = true;
-    for (auto Param : *FD->getParameters()) {
+
+    auto *Params = FD->getParameters();
+    for (auto Param : *Params) {
       if (Param == HandlerDesc.getHandler()) {
-        /// We don't need to pass the completion handler to the async method.
-        continue;
+        /// If we're not replacing the handler with anything, drop it.
+        if (HandlerReplacement.empty())
+          continue;
+
+        // If this is the last param, and we can use a trailing closure, do so.
+        if (CanUseTrailingClosure && Param == Params->back()) {
+          OS << tok::r_paren << " ";
+          OS << HandlerReplacement;
+          return;
+        }
+        // Otherwise fall through to do the replacement.
       }
-      if (!FirstParam) {
+
+      if (Param != Params->front())
         OS << tok::comma << " ";
-      } else {
-        FirstParam = false;
-      }
-      if (!Param->getArgumentName().empty()) {
+
+      if (!Param->getArgumentName().empty())
         OS << Param->getArgumentName() << tok::colon << " ";
+
+      if (Param == HandlerDesc.getHandler()) {
+        OS << HandlerReplacement;
+      } else {
+        OS << Param->getParameterName();
       }
-      OS << Param->getParameterName();
     }
     OS << tok::r_paren;
   }
@@ -6408,19 +6591,10 @@ class AsyncConverter : private SourceEntityWalker {
   /// Adds the result type of a refactored async function that previously
   /// returned results via a completion handler described by \p HandlerDesc.
   void addAsyncFuncReturnType(const AsyncHandlerDesc &HandlerDesc) {
+    // Type or (Type1, Type2, ...)
     SmallVector<Type, 2> Scratch;
-    auto ReturnTypes = HandlerDesc.getAsyncReturnTypes(Scratch);
-    if (ReturnTypes.size() > 1) {
-      OS << tok::l_paren;
-    }
-
-    llvm::interleave(
-        ReturnTypes, [&](Type Ty) { Ty->print(OS); },
-        [&]() { OS << tok::comma << " "; });
-
-    if (ReturnTypes.size() > 1) {
-      OS << tok::r_paren;
-    }
+    addTupleOf(HandlerDesc.getAsyncReturnTypes(Scratch), OS,
+               [&](auto Ty) { Ty->print(OS); });
   }
 
   /// If \p FD is generic, adds a type annotation with the return type of the
@@ -6450,6 +6624,24 @@ class AsyncConverter : private SourceEntityWalker {
   }
 };
 
+/// Adds an attribute to describe a completion handler function's async
+/// alternative if necessary.
+void addCompletionHandlerAsyncAttrIfNeccessary(
+    ASTContext &Ctx, const FuncDecl *FD,
+    const AsyncHandlerParamDesc &HandlerDesc,
+    SourceEditConsumer &EditConsumer) {
+  if (!Ctx.LangOpts.EnableExperimentalConcurrency)
+    return;
+
+  llvm::SmallString<0> HandlerAttribute;
+  llvm::raw_svector_ostream OS(HandlerAttribute);
+  OS << "@completionHandlerAsync(\"";
+  HandlerDesc.printAsyncFunctionName(OS);
+  OS << "\", completionHandlerIndex: " << HandlerDesc.Index << ")\n";
+  EditConsumer.accept(Ctx.SourceMgr, FD->getAttributeInsertionLoc(false),
+                      HandlerAttribute);
+}
+
 } // namespace asyncrefactorings
 
 bool RefactoringActionConvertCallToAsyncAlternative::isApplicable(
@@ -6571,16 +6763,7 @@ bool RefactoringActionAddAsyncAlternative::performChange() {
                       "@available(*, deprecated, message: \"Prefer async "
                       "alternative instead\")\n");
 
-  if (Ctx.LangOpts.EnableExperimentalConcurrency) {
-    // Add an attribute to describe its async alternative
-    llvm::SmallString<0> HandlerAttribute;
-    llvm::raw_svector_ostream OS(HandlerAttribute);
-    OS << "@completionHandlerAsync(\"";
-    HandlerDesc.printAsyncFunctionName(OS);
-    OS << "\", completionHandlerIndex: " << HandlerDesc.Index << ")\n";
-    EditConsumer.accept(SM, FD->getAttributeInsertionLoc(false),
-                        HandlerAttribute);
-  }
+  addCompletionHandlerAsyncAttrIfNeccessary(Ctx, FD, HandlerDesc, EditConsumer);
 
   AsyncConverter LegacyBodyCreator(TheFile, SM, DiagEngine, FD, HandlerDesc);
   if (LegacyBodyCreator.createLegacyBody()) {
@@ -6592,6 +6775,43 @@ bool RefactoringActionAddAsyncAlternative::performChange() {
 
   return false;
 }
+
+bool RefactoringActionAddAsyncWrapper::isApplicable(
+    const ResolvedCursorInfo &CursorInfo, DiagnosticEngine &Diag) {
+  using namespace asyncrefactorings;
+
+  auto *FD = findFunction(CursorInfo);
+  if (!FD)
+    return false;
+
+  auto HandlerDesc =
+      AsyncHandlerParamDesc::find(FD, /*RequireAttributeOrName=*/false);
+  return HandlerDesc.isValid();
+}
+
+bool RefactoringActionAddAsyncWrapper::performChange() {
+  using namespace asyncrefactorings;
+
+  auto *FD = findFunction(CursorInfo);
+  assert(FD &&
+         "Should not run performChange when refactoring is not applicable");
+
+  auto HandlerDesc =
+      AsyncHandlerParamDesc::find(FD, /*RequireAttributeOrName=*/false);
+  assert(HandlerDesc.isValid() &&
+         "Should not run performChange when refactoring is not applicable");
+
+  AsyncConverter Converter(TheFile, SM, DiagEngine, FD, HandlerDesc);
+  if (!Converter.createAsyncWrapper())
+    return true;
+
+  addCompletionHandlerAsyncAttrIfNeccessary(Ctx, FD, HandlerDesc, EditConsumer);
+
+  // Add the async wrapper.
+  Converter.insertAfter(FD, EditConsumer);
+  return false;
+}
+
 } // end of anonymous namespace
 
 StringRef swift::ide::