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

Author: asl

lib/AST/Decl.cpp

@@ -1391,8 +1391,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();
@@ -1405,13 +1405,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();
@@ -1428,8 +1429,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;
     }
@@ -11441,8 +11442,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());
     }
   }
@@ -11470,8 +11470,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/SIL/IR/SILFunctionType.cpp

@@ -2514,7 +2514,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());
     auto &ctx = origFd->getASTContext();
     if (auto accessor = dyn_cast<AccessorDecl>(origFd)) {
@@ -2546,6 +2546,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 = [&]{
@@ -2569,8 +2582,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

@@ -6475,17 +6475,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

@@ -2353,9 +2353,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                             bool canUnwind, 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

@@ -5813,6 +5813,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;
@@ -7245,9 +7248,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;
   }
 
@@ -7737,9 +7740,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

@@ -281,11 +281,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

@@ -7437,7 +7437,6 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
 
     case TypeKind::Error:
     case TypeKind::Unresolved:
-    case TypeKind::YieldResult:
       return getTypeMatchFailure(locator);
 
     case TypeKind::BuiltinFixedArray: {
@@ -7456,6 +7455,24 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
                         locator.withPathElement(
                             LocatorPathElt::GenericArgument(0)));
     }
+      
+    case TypeKind::YieldResult: {
+      if (kind != ConstraintKind::Bind && kind != ConstraintKind::Subtype)
+        return getTypeMatchFailure(locator);
+
+      auto *yield1 = cast<YieldResultType>(desugar1);
+      auto *yield2 = cast<YieldResultType>(desugar2);
+
+      // TODO: In theory we can convert inout yield to non-inout one,
+      // however, we disallow this for now as overall generic coroutine
+      // semantics is a bit vague.
+      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
@@ -8469,7 +8486,6 @@ ConstraintSystem::simplifyConstructionConstraint(
   case TypeKind::Unresolved:
   case TypeKind::Error:
   case TypeKind::Placeholder:
-  case TypeKind::YieldResult:
     return SolutionKind::Error;
 
   case TypeKind::GenericFunction:
@@ -8569,6 +8585,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

@@ -928,8 +928,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 {

test/SILGen/witnesses.swift

@@ -464,8 +464,7 @@ class PropertyRequirementWitnessFromBase : PropertyRequirementBase, PropertyRequ
   // CHECK-NEXT: [[METH:%.*]] = class_method [[CAST_ARG2_LOADED]] : $PropertyRequirementBase, #PropertyRequirementBase.width!modify
   // CHECK-NEXT: ([[RES:%.*]], [[TOKEN:%.*]]) = begin_apply [[METH]]([[CAST_ARG2_LOADED]]) : $@yield_once @convention(method) (@guaranteed PropertyRequirementBase) -> @yields @inout Int
   // CHECK-NEXT: yield [[RES]]
-  // CHECK:      end_apply [[TOKEN]]
-  // CHECK-NEXT: [[TUPLE:%.*]] = tuple ()
+  // CHECK:      [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
   // CHECK-NEXT: end_borrow [[ARG2_LOADED]]
   // CHECK-NEXT: return [[TUPLE]]
 
@@ -474,8 +473,7 @@ class PropertyRequirementWitnessFromBase : PropertyRequirementBase, PropertyRequ
   // CHECK: [[METH:%.*]] = function_ref @$s9witnesses23PropertyRequirementBaseC6heightSivMZ
   // CHECK-NEXT: ([[RES:%.*]], [[TOKEN:%.*]]) = begin_apply [[METH]]
   // CHECK-NEXT: yield [[RES]]
-  // CHECK:      end_apply [[TOKEN]]
-  // CHECK-NEXT: [[TUPLE:%.*]] = tuple ()
+  // CHECK:      [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
   // CHECK-NEXT: return [[TUPLE]]
 
  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s9witnesses34PropertyRequirementWitnessFromBaseCAA0bC0A2aDP5depthSivMTW
@@ -484,8 +482,7 @@ class PropertyRequirementWitnessFromBase : PropertyRequirementBase, PropertyRequ
   // CHECK: [[METH:%.*]] = class_method [[ARG2_LOADED]] : $PropertyRequirementWitnessFromBase, #PropertyRequirementWitnessFromBase.depth!modify
   // CHECK-NEXT: ([[RES:%.*]], [[TOKEN:%.*]]) = begin_apply [[METH]]([[ARG2_LOADED]]) : $@yield_once @convention(method) (@guaranteed PropertyRequirementWitnessFromBase) -> @yields @inout Int
   // CHECK-NEXT: yield [[RES]]
-  // CHECK:      end_apply [[TOKEN]]
-  // CHECK-NEXT: [[TUPLE:%.*]] = tuple ()
+  // CHECK:      [[TUPLE:%.*]] = end_apply [[TOKEN]] as $()
   // CHECK-NEXT: end_borrow [[ARG2_LOADED]]
   // CHECK-NEXT: return [[TUPLE]]
 }