First cut of making coroutine AST type

Author: asl

include/swift/AST/AnyFunctionRef.h

@@ -96,7 +96,7 @@ class AnyFunctionRef {
   Type getBodyResultType() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
       if (auto *FD = dyn_cast<FuncDecl>(AFD))
-        return FD->mapTypeIntoContext(FD->getResultInterfaceType());
+        return FD->mapTypeIntoContext(FD->getResultInterfaceTypeWithoutYields());
       return TupleType::getEmpty(AFD->getASTContext());
     }
     return cast<AbstractClosureExpr *>(TheFunction)->getResultType();

include/swift/AST/Decl.h

@@ -8508,9 +8508,15 @@ class FuncDecl : public AbstractFunctionDecl {
     return FnRetType.getSourceRange();
   }
 
-  /// Retrieve the result interface type of this function.
+  /// Retrieve the full result interface type of this function, including yields
   Type getResultInterfaceType() const;
 
+  /// Same as above, but without @yields
+  Type getResultInterfaceTypeWithoutYields() const;
+
+  /// Same as above, but only yields
+  Type getYieldsInterfaceType() const;
+
   /// Returns the result interface type of this function if it has already been
   /// computed, otherwise `nullopt`. This should only be used for dumping.
   std::optional<Type> getCachedResultInterfaceType() const;

include/swift/AST/Types.h

@@ -1724,10 +1724,13 @@ class YieldResultType : public TypeBase {
 };
 
 BEGIN_CAN_TYPE_WRAPPER(YieldResultType, Type)
-PROXY_CAN_TYPE_SIMPLE_GETTER(getResultType)
-static CanYieldResultType get(CanType type, bool InOut) {
-  return CanYieldResultType(YieldResultType::get(type, InOut));
-}
+  PROXY_CAN_TYPE_SIMPLE_GETTER(getResultType)
+  bool isInOut() const {
+    return getPointer()->isInOut();
+  }
+  static CanYieldResultType get(CanType type, bool InOut) {
+    return CanYieldResultType(YieldResultType::get(type, InOut));
+  }
 END_CAN_TYPE_WRAPPER(YieldResultType, Type)
 
 /// BuiltinType - An abstract class for all the builtin types.

include/swift/SIL/AbstractionPattern.h

@@ -1524,6 +1524,10 @@ class AbstractionPattern {
   /// type, return the abstraction pattern for one of its argument types.
   AbstractionPattern getParameterizedProtocolArgType(unsigned i) const;
 
+  /// Given that the value being abstracted is a yield result type,
+  /// return the abstraction pattern for corresponding type.
+  AbstractionPattern getYieldResultType() const;
+
   /// Given that the value being abstracted is a function, return the
   /// abstraction pattern for its result type.
   AbstractionPattern getFunctionResultType() const;

lib/AST/ASTVerifier.cpp

@@ -1098,7 +1098,7 @@ class Verifier : public ASTWalker {
       auto func = Functions.back();
       Type resultType;
       if (auto *FD = dyn_cast<FuncDecl>(func)) {
-        resultType = FD->getResultInterfaceType();
+        resultType = FD->getResultInterfaceTypeWithoutYields();
         resultType = FD->mapTypeIntoContext(resultType);
       } else if (auto closure = dyn_cast<AbstractClosureExpr>(func)) {
         resultType = closure->getResultType();

lib/AST/Decl.cpp

@@ -1374,6 +1374,18 @@ bool AbstractFunctionDecl::isTransparent() const {
   return false;
 }
 
+bool AbstractFunctionDecl::isCoroutine() const {
+  // Check if the declaration had the attribute.
+  if (getAttrs().hasAttribute<CoroutineAttr>())
+    return true;
+
+  // If this is an accessor, then check if its a coroutine.
+  if (const auto *AD = dyn_cast<AccessorDecl>(this))
+    return AD->isCoroutine();
+
+  return false;
+}
+
 bool ParameterList::hasInternalParameter(StringRef Prefix) const {
   for (auto param : *this) {
     if (param->hasName() && param->getNameStr().starts_with(Prefix))
@@ -11352,6 +11364,67 @@ std::optional<Type> FuncDecl::getCachedResultInterfaceType() const {
   return ResultTypeRequest{mutableThis}.getCachedResult();
 }
 
+Type FuncDecl::getResultInterfaceTypeWithoutYields() const {
+  auto resultType = getResultInterfaceType();
+  if (resultType->hasError())
+    return resultType;
+
+  // Coroutine result type should either be a yield result
+  // or a tuple containing both yielded and normal result types.
+  // In both cases, strip the @yield result types
+  if (isCoroutine()) {
+    if (auto *tupleResTy = resultType->getAs<TupleType>()) {
+      // Strip @yield results on the first level of tuple
+      SmallVector<TupleTypeElt, 4> elements;
+      for (const auto &elt : tupleResTy->getElements()) {
+        Type eltTy = elt.getType();
+        if (eltTy->is<YieldResultType>())
+          continue;
+        elements.push_back(eltTy);
+      }
+
+      // Handle vanishing tuples --  flatten to produce the
+      // element type.
+      if (elements.size() == 1)
+          resultType = elements[0].getType();
+      else
+          resultType = TupleType::get(elements, getASTContext());
+    } else {
+      assert(resultType->is<YieldResultType>());
+      resultType = TupleType::getEmpty(getASTContext());
+    }
+  }
+
+  return resultType;
+}
+
+Type FuncDecl::getYieldsInterfaceType() const {
+  auto resultType = getResultInterfaceType();
+  if (resultType->hasError())
+    return resultType;
+
+  if (!isCoroutine())
+    return TupleType::getEmpty(getASTContext());
+
+  // Coroutine result type should either be a yield result
+  // or a tuple containing both yielded and normal result types.
+  // In both cases, strip the @yield result types
+  if (auto *tupleResTy = resultType->getAs<TupleType>()) {
+    // Keep @yield results on the first level of tuple
+    for (const auto &elt : tupleResTy->getElements()) {
+      Type eltTy = elt.getType();
+      if (eltTy->is<YieldResultType>())
+        return eltTy;
+      }
+
+    llvm_unreachable("coroutine must have a yield result");
+  } else {
+    assert(resultType->is<YieldResultType>());
+  }
+
+  return resultType;
+}
+
 bool FuncDecl::isUnaryOperator() const {
   if (!isOperator())
     return false;

lib/SIL/IR/AbstractionPattern.cpp

@@ -1038,6 +1038,13 @@ AbstractionPattern::getParameterizedProtocolArgType(unsigned argIndex) const {
             cast<ParameterizedProtocolType>(getType()).getArgs()[argIndex]);
 }
 
+AbstractionPattern AbstractionPattern::getYieldResultType() const {
+  assert(getKind() == Kind::Type);
+  return AbstractionPattern(getGenericSubstitutions(),
+                            getGenericSignature(),
+                            cast<YieldResultType>(getType()).getResultType());
+}
+
 AbstractionPattern AbstractionPattern::removingMoveOnlyWrapper() const {
   switch (getKind()) {
   case Kind::Invalid:
@@ -2731,6 +2738,13 @@ class SubstFunctionTypePatternVisitor
     llvm_unreachable("shouldn't encounter pack element by itself");
   }
 
+  CanType visitYieldResultType(CanYieldResultType yield,
+                               AbstractionPattern pattern) {
+    auto resultType = visit(yield.getResultType(), pattern.getYieldResultType());
+    return YieldResultType::get(resultType, yield->isInOut())
+      ->getCanonicalType();
+  }
+
   CanType handlePackExpansion(AbstractionPattern origExpansion,
                               CanType candidateSubstType) {
     // When we're within a pack expansion, pack references matching that

lib/SIL/IR/SILFunctionType.cpp

@@ -1417,6 +1417,10 @@ class DestructureResults {
       return;
     }
 
+    // Skip yields, they should've already processed elsewhere
+    if (isa<YieldResultType>(substType))
+      return;
+
     auto &substResultTLForConvention = TC.getTypeLowering(
         origType, substType, TypeExpansionContext::minimal());
     auto &substResultTL = TC.getTypeLowering(origType, substType,
@@ -2317,23 +2321,22 @@ lowerCaptureContextParameters(TypeConverter &TC, SILDeclRef function,
          "iterating over loweredCaptures.getCaptures().");
 }
 
-static AccessorDecl *
-getAsCoroutineAccessor(std::optional<SILDeclRef> constant) {
+static FuncDecl *getAsCoroutine(std::optional<SILDeclRef> constant) {
   if (!constant || !constant->hasDecl())
     return nullptr;;
 
-  auto accessor = dyn_cast<AccessorDecl>(constant->getDecl());
-  if (!accessor || !accessor->isCoroutine())
+  auto fd = dyn_cast<FuncDecl>(constant->getDecl());
+  if (!fd || !fd->isCoroutine())
     return nullptr;
 
-  return accessor;
+  return fd;
 }
 
 static void destructureYieldsForReadAccessor(TypeConverter &TC,
-                                         TypeExpansionContext expansion,
-                                         AbstractionPattern origType,
-                                         CanType valueType,
-                                         SmallVectorImpl<SILYieldInfo> &yields){
+                                             TypeExpansionContext expansion,
+                                             AbstractionPattern origType,
+                                             CanType valueType,
+                                             SmallVectorImpl<SILYieldInfo> &yields){
   // Recursively destructure tuples.
   if (origType.isTuple()) {
     auto valueTupleType = cast<TupleType>(valueType);
@@ -2365,25 +2368,19 @@ static void destructureYieldsForReadAccessor(TypeConverter &TC,
 
 static void destructureYieldsForCoroutine(TypeConverter &TC,
                                           TypeExpansionContext expansion,
-                                          std::optional<SILDeclRef> constant,
                                           AbstractionPattern origType,
                                           CanType canValueType,
-                                          SmallVectorImpl<SILYieldInfo> &yields,
-                                          SILCoroutineKind &coroutineKind) {
-  auto accessor = getAsCoroutineAccessor(constant);
-  if (!accessor)
-    return;
-
+                                          bool isInOutYield,
+                                          SmallVectorImpl<SILYieldInfo> &yields) {
   // 'modify' yields an inout of the target type.
-  if (isYieldingMutableAccessor(accessor->getAccessorKind())) {
+  if (isInOutYield) {
     auto loweredValueTy =
         TC.getLoweredType(origType, canValueType, expansion);
     yields.push_back(SILYieldInfo(loweredValueTy.getASTType(),
                                   ParameterConvention::Indirect_Inout));
   } else {
     // 'read' yields a borrowed value of the target type, destructuring
     // tuples as necessary.
-    assert(isYieldingImmutableAccessor(accessor->getAccessorKind()));
     destructureYieldsForReadAccessor(TC, expansion, origType, canValueType,
                                      yields);
   }
@@ -2511,37 +2508,44 @@ static CanSILFunctionType getSILFunctionType(
 
   bool hasSendingResult = substFnInterfaceType->getExtInfo().hasSendingResult();
 
-  // Get the yield type for an accessor coroutine.
+  // Get the yield type for coroutine.
   SILCoroutineKind coroutineKind = SILCoroutineKind::None;
   AbstractionPattern coroutineOrigYieldType = AbstractionPattern::getInvalid();
   CanType coroutineSubstYieldType;
 
-  if (auto accessor = getAsCoroutineAccessor(constant)) {
-    auto origAccessor = cast<AccessorDecl>(origConstant->getDecl());
-    auto &ctx = origAccessor->getASTContext();
-    coroutineKind =
+  bool isInOutYield = false;
+  if (auto fd = getAsCoroutine(constant)) {
+    auto origFd = cast<FuncDecl>(origConstant->getDecl());
+    auto &ctx = origFd->getASTContext();
+    if (auto accessor = dyn_cast<AccessorDecl>(origFd)) {
+      coroutineKind =
         (requiresFeatureCoroutineAccessors(accessor->getAccessorKind()) &&
          ctx.SILOpts.CoroutineAccessorsUseYieldOnce2)
-            ? SILCoroutineKind::YieldOnce2
-            : SILCoroutineKind::YieldOnce;
-
+        ? SILCoroutineKind::YieldOnce2
+        : SILCoroutineKind::YieldOnce;
+    } else {
+      // FIXME: Decide we'd directly go to YieldOnce2 for non-accessor coroutines
+      coroutineKind = SILCoroutineKind::YieldOnce;
+    }
+    
     // Coroutine accessors are always native, so fetch the native
     // abstraction pattern.
-    auto origStorage = origAccessor->getStorage();
-    coroutineOrigYieldType = TC.getAbstractionPattern(origStorage,
-                                                      /*nonobjc*/ true)
-                               .getReferenceStorageReferentType();
-
-    auto storage = accessor->getStorage();
-    auto valueType = storage->getValueInterfaceType();
+    auto sig = origFd->getGenericSignatureOfContext()
+      .getCanonicalSignature();
+    auto origYieldType = origFd->getYieldsInterfaceType()->castTo<YieldResultType>();
+    auto reducedYieldType = sig.getReducedType(origYieldType->getResultType());
+    coroutineOrigYieldType = AbstractionPattern(sig, reducedYieldType);
 
+    auto yieldType = fd->getYieldsInterfaceType()->castTo<YieldResultType>();
+    auto valueType = yieldType->getResultType();
+    isInOutYield = yieldType->isInOut();
     if (reqtSubs) {
       valueType = valueType.subst(*reqtSubs);
       coroutineSubstYieldType = valueType->getReducedType(
           genericSig);
     } else {
       coroutineSubstYieldType = valueType->getReducedType(
-          accessor->getGenericSignature());
+          fd->getGenericSignature());
     }
   }
 
@@ -2566,7 +2570,8 @@ 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 && getAsCoroutineAccessor(constant))
+    if (constant != origConstant &&
+        coroutineKind != SILCoroutineKind::None)
       return true;
 
     // We don't currently use substituted function types for generic function
@@ -2670,10 +2675,12 @@ static CanSILFunctionType getSILFunctionType(
 
   // Destructure the coroutine yields.
   SmallVector<SILYieldInfo, 8> yields;
-  destructureYieldsForCoroutine(TC, expansionContext, constant,
-                                coroutineOrigYieldType, coroutineSubstYieldType,
-                                yields, coroutineKind);
-  
+  if (coroutineKind != SILCoroutineKind::None) {
+    destructureYieldsForCoroutine(TC, expansionContext,
+                                  coroutineOrigYieldType, coroutineSubstYieldType,
+                                  isInOutYield, yields);
+  }
+
   // Destructure the result tuple type.
   SmallVector<SILResultInfo, 8> results;
   {
@@ -4977,6 +4984,8 @@ TypeConverter::getLoweredFormalTypes(SILDeclRef constant,
     extInfo = extInfo.withThrows(true, innerExtInfo.getThrownError());
   if (innerExtInfo.isAsync())
     extInfo = extInfo.withAsync(true);
+  if (innerExtInfo.isCoroutine())
+    extInfo = extInfo.withCoroutine(true);
 
   // Distributed thunks are always `async throws`
   if (constant.isDistributedThunk()) {

lib/SILGen/SILGenFunction.cpp

@@ -1121,15 +1121,15 @@ void SILGenFunction::emitFunction(FuncDecl *fd) {
 
   auto captureInfo = SGM.M.Types.getLoweredLocalCaptures(SILDeclRef(fd));
   emitProlog(fd, captureInfo, fd->getParameters(), fd->getImplicitSelfDecl(),
-             fd->getResultInterfaceType(), fd->getEffectiveThrownErrorType(),
+             fd->getResultInterfaceTypeWithoutYields(), fd->getEffectiveThrownErrorType(),
              fd->getThrowsLoc());
 
   if (fd->isDistributedActorFactory()) {
     // Synthesize the factory function body
     emitDistributedActorFactory(fd);
   } else {
     prepareEpilog(fd,
-                  fd->getResultInterfaceType(),
+                  fd->getResultInterfaceTypeWithoutYields(),
                   fd->getEffectiveThrownErrorType(),
                   CleanupLocation(fd));
 

lib/SILGen/SILGenProlog.cpp

@@ -1568,6 +1568,10 @@ static void emitIndirectResultParameters(SILGenFunction &SGF,
 
   assert(!resultType->is<PackExpansionType>());
 
+  // Skip yields, they are emitted elsewhere
+  if (resultType->is<YieldResultType>())
+    return;
+
   // If the return type is address-only, emit the indirect return argument.
 
   // The calling convention always uses minimal resilience expansion.