Initial & rudimentary support for coroutine function types. Some fixes while here

Author: asl

lib/AST/Decl.cpp

@@ -1171,8 +1171,8 @@ AnyFunctionRef::getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buf
                                     bool mapIntoContext) const {
   assert(buffer.empty());
   if (auto *AFD = getAbstractFunctionDecl()) {
+    // FIXME: AccessorDecl case is not necessary
     if (auto *AD = dyn_cast<AccessorDecl>(AFD)) {
-      // FIXME: AccessorDecl case is not necessary
       if (AD->isCoroutine()) {
         auto valueTy = AD->getStorage()->getValueInterfaceType()
                                        ->getReferenceStorageReferent();
@@ -1185,13 +1185,14 @@ AnyFunctionRef::getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buf
         return buffer;
       }
     } else if (AFD->isCoroutine()) {
-      auto resType = AFD->getInterfaceType()->castTo<FunctionType>()->getResult();
-      if (auto *resFnType = resType->getAs<FunctionType>())
-        resType = resFnType->getResult();
-
-      if (resType->hasError())
+      auto fnType = AFD->getInterfaceType();
+      if (fnType->hasError())
         return {};
 
+      auto resType = fnType->castTo<AnyFunctionType>()->getResult();
+      if (auto *resFnType = resType->getAs<AnyFunctionType>())
+        resType = resFnType->getResult();
+
       auto addYieldInfo =
         [&](const YieldResultType *yieldResultTy) {
           Type valueTy = yieldResultTy->getResultType();
@@ -1208,8 +1209,8 @@ AnyFunctionRef::getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buf
           if (auto *yieldResTy = eltTy->getAs<YieldResultType>())
             addYieldInfo(yieldResTy);
         }
-      else
-        addYieldInfo(resType->castTo<YieldResultType>());
+      else if (auto *yieldResTy = resType->getAs<YieldResultType>())
+        addYieldInfo(yieldResTy);
 
       return buffer;
     }
@@ -10822,8 +10823,7 @@ Type FuncDecl::getResultInterfaceTypeWithoutYields() const {
           resultType = elements[0].getType();
       else
           resultType = TupleType::get(elements, getASTContext());
-    } else {
-      assert(resultType->is<YieldResultType>());
+    } else if (resultType->is<YieldResultType>()) {
       resultType = TupleType::getEmpty(getASTContext());
     }
   }
@@ -10851,8 +10851,8 @@ Type FuncDecl::getYieldsInterfaceType() const {
       }
 
     llvm_unreachable("coroutine must have a yield result");
-  } else {
-    assert(resultType->is<YieldResultType>());
+  } else if (!resultType->is<YieldResultType>()) {
+    resultType = TupleType::getEmpty(getASTContext());
   }
 
   return resultType;

lib/IRGen/IRGenSIL.cpp

@@ -4787,13 +4787,13 @@ void IRGenSILFunction::visitEndApply(BeginApplyInst *i, EndApplyInst *ei) {
 
   if (!isAbort) {
     auto resultType = call->getType();
+    Explosion e;
     if (!resultType->isVoidTy()) {
-      Explosion e;
       // FIXME: Do we need to handle ABI-related conversions here?
       // It seems we cannot have C function convention for coroutines, etc.
       extractScalarResults(*this, resultType, call, e);
-      setLoweredExplosion(ei, e);
     }
+    setLoweredExplosion(ei, e);
   }
 
   coroutine.Temporaries.destroyAll(*this);

lib/SIL/IR/SILFunctionType.cpp

@@ -2380,7 +2380,7 @@ static CanSILFunctionType getSILFunctionType(
   CanType coroutineSubstYieldType;
 
   bool isInOutYield = false;
-  if (auto fd = getAsCoroutine(constant)) {
+  if (auto fd = getAsCoroutine(constant)) { // Derive yield type for declaration
     auto origFd = cast<FuncDecl>(origConstant->getDecl());
     if (auto accessor = dyn_cast<AccessorDecl>(origFd)) {
       coroutineKind =
@@ -2410,6 +2410,19 @@ static CanSILFunctionType getSILFunctionType(
       coroutineSubstYieldType = valueType->getReducedType(
           fd->getGenericSignature());
     }
+  } else if (substFnInterfaceType->isCoroutine()) { // Derive yield type for function type
+    coroutineKind = SILCoroutineKind::YieldOnce;
+    auto origYieldType = origType.getFunctionResultType().getType()->castTo<YieldResultType>();
+    auto reducedYieldType = genericSig.getReducedType(origYieldType->getResultType());
+    coroutineOrigYieldType = AbstractionPattern(genericSig, reducedYieldType);
+
+    auto yieldType = substFnInterfaceType->getResult()->castTo<YieldResultType>();
+    auto valueType = yieldType->getResultType();
+    isInOutYield = yieldType->isInOut();
+    if (reqtSubs)
+      valueType = valueType.subst(*reqtSubs);
+
+    coroutineSubstYieldType = valueType->getReducedType(genericSig);
   }
 
   bool shouldBuildSubstFunctionType = [&]{
@@ -2433,8 +2446,7 @@ static CanSILFunctionType getSILFunctionType(
     // for class override thunks.  This is required to make the yields
     // match in abstraction to the base method's yields, which is necessary
     // to make the extracted continuation-function signatures match.
-    if (constant != origConstant &&
-        coroutineKind != SILCoroutineKind::None)
+    if (constant != origConstant && getAsCoroutine(constant))
       return true;
 
     // We don't currently use substituted function types for generic function

lib/SILGen/SILGenApply.cpp

@@ -6115,17 +6115,19 @@ SILGenFunction::emitBeginApplyWithRethrow(SILLocation loc, SILValue fn,
   return {token, abortCleanup, allocation, deallocCleanup};
 }
 
-void SILGenFunction::emitEndApplyWithRethrow(
+SILValue SILGenFunction::emitEndApplyWithRethrow(
     SILLocation loc, MultipleValueInstructionResult *token,
-    SILValue allocation) {
+    SILValue allocation,
+    SILType resultType) {
   // TODO: adjust this to handle results of TryBeginApplyInst.
   assert(token->isBeginApplyToken());
 
-  B.createEndApply(loc, token,
-                   SILType::getEmptyTupleType(getASTContext()));
+  SILValue result =
+    B.createEndApply(loc, token, resultType);
   if (allocation) {
     B.createDeallocStack(loc, allocation);
   }
+  return result;
 }
 
 void SILGenFunction::emitYield(SILLocation loc,

lib/SILGen/SILGenFunction.h

@@ -2196,9 +2196,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   emitBeginApplyWithRethrow(SILLocation loc, SILValue fn, SILType substFnType,
                             SubstitutionMap subs, ArrayRef<SILValue> args,
                             SmallVectorImpl<SILValue> &yields);
-  void emitEndApplyWithRethrow(SILLocation loc,
-                               MultipleValueInstructionResult *token,
-                               SILValue allocation);
+  SILValue emitEndApplyWithRethrow(SILLocation loc,
+                                   MultipleValueInstructionResult *token,
+                                   SILValue allocation,
+                                   SILType resultType);
 
   ManagedValue emitExtractFunctionIsolation(SILLocation loc,
                                         ArgumentSource &&fnValue);

lib/SILGen/SILGenPoly.cpp

@@ -5541,6 +5541,9 @@ static ManagedValue createThunk(SILGenFunction &SGF,
   assert(expectedType->getLanguage() ==
          fn.getType().castTo<SILFunctionType>()->getLanguage() &&
          "bridging in re-abstraction thunk?");
+  // We cannot reabstract coroutines (yet)
+  assert(!expectedType->isCoroutine() && !sourceType->isCoroutine() &&
+         "cannot reabstract a coroutine");
 
   // Declare the thunk.
   SubstitutionMap interfaceSubs;
@@ -5554,6 +5557,7 @@ static ManagedValue createThunk(SILGenFunction &SGF,
                                       genericEnv,
                                       interfaceSubs,
                                       dynamicSelfType);
+
   // An actor-isolated non-async function can be converted to an async function
   // by inserting a hop to the global actor.
   CanType globalActorForThunk;
@@ -6849,9 +6853,9 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
     }
 
     // End the inner coroutine normally.
-    emitEndApplyWithRethrow(loc, token, allocation);
+    result = emitEndApplyWithRethrow(loc, token, allocation,
+                                     SILType::getEmptyTupleType(getASTContext()));
 
-    result = B.createTuple(loc, {});
     break;
   }
 
@@ -7241,9 +7245,8 @@ void SILGenFunction::emitProtocolWitness(
     }
 
     // End the inner coroutine normally.
-    emitEndApplyWithRethrow(loc, token, allocation);
-
-    reqtResultValue = B.createTuple(loc, {});
+    reqtResultValue = emitEndApplyWithRethrow(loc, token, allocation,
+                                              SILType::getEmptyTupleType(getASTContext()));
     break;
   }
 

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -255,11 +255,22 @@ class BeginApplySite {
       for (auto calleeYield : BeginApply->getYieldedValues()) {
         calleeYield->replaceAllUsesWith(SILUndef::get(calleeYield));
       }
+
+      if (EndApply)
+        EndApply->replaceAllUsesWith(SILUndef::get(EndApply));
     }
 
     // Remove the resumption sites.
-    if (EndApply)
+    if (EndApply) {
+      // All potential users of end_apply should've been replaced above. The only
+      // case where we might end with more users is when end_apply itself is
+      // unreachable. Make sure the function is well-formed and replace the
+      // results with undef.
+      if (!EndApply->use_empty())
+        EndApply->replaceAllUsesWith(SILUndef::get(EndApply));
+
       EndApply->eraseFromParent();
+    }
     if (AbortApply)
       AbortApply->eraseFromParent();
     for (auto *EndBorrow : EndBorrows)

lib/Sema/CSSimplify.cpp

@@ -7211,9 +7211,26 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
 
     case TypeKind::Error:
     case TypeKind::Unresolved:
-    case TypeKind::YieldResult:
       return getTypeMatchFailure(locator);
 
+    case TypeKind::YieldResult: {
+      if (simplifyType(desugar1)->isEqual(simplifyType(desugar2)))
+        return getTypeMatchSuccess();
+
+      if (kind != ConstraintKind::Bind)
+        return getTypeMatchFailure(locator);
+
+      auto *yield1 = cast<YieldResultType>(desugar1);
+      auto *yield2 = cast<YieldResultType>(desugar2);
+
+      if (yield1->isInOut() != yield2->isInOut())
+        return getTypeMatchFailure(locator);
+
+      return matchTypes(yield1->getResultType(), yield2->getResultType(),
+                        ConstraintKind::Bind, subflags,
+                  locator.withPathElement(ConstraintLocator::LValueConversion));
+    }
+
     case TypeKind::Placeholder: {
       // If it's allowed to attempt fixes, let's delegate
       // decision to `repairFailures`, since depending on
@@ -8213,7 +8230,6 @@ ConstraintSystem::simplifyConstructionConstraint(
   case TypeKind::Unresolved:
   case TypeKind::Error:
   case TypeKind::Placeholder:
-  case TypeKind::YieldResult:
     return SolutionKind::Error;
 
   case TypeKind::GenericFunction:
@@ -8313,6 +8329,7 @@ ConstraintSystem::simplifyConstructionConstraint(
   case TypeKind::Function:
   case TypeKind::LValue:
   case TypeKind::InOut:
+  case TypeKind::YieldResult:
   case TypeKind::Module:
   case TypeKind::Pack:
   case TypeKind::PackExpansion:

lib/Sema/TypeCheckGeneric.cpp

@@ -777,8 +777,12 @@ GenericSignatureRequest::evaluate(Evaluator &evaluator,
         }
       }();
       if (resultTypeRepr && !resultTypeRepr->hasOpaque()) {
+        bool isCoroutine = func ? func->isCoroutine() : false;
+        TypeResolutionOptions resultOptions(TypeResolverContext::FunctionResult);
+        if (isCoroutine)
+          resultOptions |= TypeResolutionFlags::Coroutine;
         const auto resultType =
-            resolution.withOptions(TypeResolverContext::FunctionResult)
+            resolution.withOptions(resultOptions)
                 .resolveType(resultTypeRepr);
 
         inferenceSources.push_back(resultType.getPointer());

test/SILGen/modify.swift

@@ -80,10 +80,9 @@ extension Derived : Abstractable {}
 // CHECK-NEXT: [[T1:%.*]] = partial_apply [callee_guaranteed] [[REABSTRACTOR]]([[CVT_FN]])
 // CHECK-NEXT: store [[T1]] to [init] [[SUPER_ADDR]]
 // CHECK-NEXT: dealloc_stack [[SUB_ADDR]]
-// CHECK-NEXT: end_apply [[TOKEN]]
-// CHECK-NEXT: tuple ()
+// CHECK-NEXT: [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
 // CHECK-NEXT: end_borrow [[T0]]
-// CHECK-NEXT: return
+// CHECK-NEXT: return [[TUPLE]]
 
 // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s6modify7DerivedCAA12AbstractableA2aDP19finalStoredFunction6ResultQzycvMTW
 // CHECK: bb0(%0 : $*Derived):
@@ -108,10 +107,9 @@ extension Derived : Abstractable {}
 // CHECK-NEXT: [[T1:%.*]] = partial_apply [callee_guaranteed] [[REABSTRACTOR]]([[CVT_FN]])
 // CHECK-NEXT: store [[T1]] to [init] [[SUPER_ADDR]]
 // CHECK-NEXT: dealloc_stack [[SUB_ADDR]]
-// CHECK-NEXT: end_apply [[TOKEN]]
-// CHECK-NEXT: tuple ()
+// CHECK-NEXT: [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
 // CHECK-NEXT: end_borrow [[T0]]
-// CHECK-NEXT: return
+// CHECK-NEXT: return [[TUPLE]]
 
 // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s6modify7DerivedCAA12AbstractableA2aDP14staticFunction6ResultQzycvMZTW
 // CHECK: bb0(%0 : $@thick Derived.Type):
@@ -135,9 +133,8 @@ extension Derived : Abstractable {}
 // CHECK-NEXT: [[T1:%.*]] = partial_apply [callee_guaranteed] [[REABSTRACTOR]]([[CVT_FN]])
 // CHECK-NEXT: store [[T1]] to [init] [[SUPER_ADDR]]
 // CHECK-NEXT: dealloc_stack [[SUB_ADDR]]
-// CHECK-NEXT: end_apply [[TOKEN]]
-// CHECK-NEXT: tuple ()
-// CHECK-NEXT: return
+// CHECK-NEXT: [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
+// CHECK-NEXT: return [[TUPLE]]
 
 protocol ClassAbstractable : class {
   associatedtype Result
@@ -298,8 +295,7 @@ struct Bill : Totalled {
 // CHECK-NEXT:     ([[T1:%.*]], [[TOKEN:%.*]]) = begin_apply [[T0]]([[SELF]])
 // CHECK-NEXT:     yield [[T1]] : $*Int, resume bb1, unwind bb2
 // CHECK:        bb1:
-// CHECK-NEXT:     end_apply [[TOKEN]]
-// CHECK-NEXT:     [[T1:%.*]] = tuple ()
+// CHECK-NEXT:     [[T1:%.*]] = end_apply [[TOKEN]] as $()
 // CHECK-NEXT:     return [[T1]] :
 
 protocol AddressOnlySubscript {