Add partial_apply support for coroutines (#71653)

The patch adds lowering of partial_apply instructions for coroutines.

This pattern seems to trigger a lot of type mismatch errors in IRGen, because
coroutine functions are not substituted in the same way as regular functions
(see the patch 07f03bd "Use pattern substitutions to consistently abstract
yields" for more details). 

Other than that, lowering of partial_apply for coroutines is straightforward: we
generate another coroutine that captures arguments passed to the partial_apply
instructions. It calls the original coroutine for yields (first return) and
yields the resulting values. Then it calls the original function's continuation
for return or unwind, and forwards them to the caller as well.

After IRGen, LLVM's Coroutine pass transforms the generated coroutine (along with
all other coroutines) and eliminates llvm.coro.* intrinsics. LIT tests check
LLVM IR after this transformation.

Co-authored-by: Anton Korobeynikov <[email protected]>
Co-authored-by: Arnold Schwaighofer <[email protected]>

Author: 3 people

docs/SIL.rst

@@ -804,7 +804,7 @@ the ``@yields`` attribute.  A yielded value may have a convention attribute,
 taken from the set of parameter attributes and interpreted as if the yield
 site were calling back to the calling function.
 
-Currently, a coroutine may not have normal results.
+In addition to yielded values a coroutine could also have normal results.
 
 Coroutine functions may be used in many of the same ways as normal
 function values.  However, they cannot be called with the standard
@@ -6330,6 +6330,16 @@ callee function (and thus said signature). Instead:
    ``@inout_aliasable`` parameter convention is used when a ``@noescape``
    closure captures an ``inout`` argument.
 
+**Coroutines** ``partial_apply`` could be used to create closures over
+coroutines. Overall, the ``partial_apply`` of a coroutine is straightforward: it
+is another coroutine that captures arguments passed to the ``partial_apply``
+instruction. This closure applies the original coroutine (similar to the
+``begin_apply`` instruction) for yields (suspend) and yields the resulting
+values. Then it calls the original coroutine continuation for return or unwind,
+and forwards the results (if any) to the caller as well. Currently only the
+autodiff transformation produces ``partial_apply`` for coroutines while
+differentiating modify accessors.
+
 **NOTE:** If the callee is generic, all of its generic parameters must be bound
 by the given substitution list. The arguments are given with these generic
 substitutions applied, and the resulting closure is of concrete function type

lib/IRGen/GenDecl.cpp

@@ -6163,6 +6163,7 @@ IRGenModule::getAddrOfContinuationPrototype(CanSILFunctionType fnType) {
   llvm::Function *&entry = GlobalFuncs[entity];
   if (entry) return entry;
 
+  GenericContextScope scope(*this, fnType->getInvocationGenericSignature());
   auto signature = Signature::forCoroutineContinuation(*this, fnType);
   LinkInfo link = LinkInfo::get(*this, entity, NotForDefinition);
   entry = createFunction(*this, link, signature);