Add AST representation for coroutines

include/swift/ABI/MetadataValues.h

@@ -1296,7 +1296,10 @@ class TargetExtendedFunctionTypeFlags {
     NonIsolatedNonsending  = 0x00000004U,
 
     // Values if we have a sending result.
-    HasSendingResult  = 0x00000010U,
+    HasSendingResult       = 0x00000010U,
+
+    // Values if we have any yields
+    IsCoroutine            = 0x00000020U,
 
     /// A InvertibleProtocolSet in the high bits.
     InvertedProtocolshift = 16,
@@ -1338,6 +1341,12 @@ class TargetExtendedFunctionTypeFlags {
         (newValue ? HasSendingResult : 0));
   }
 
+  const TargetExtendedFunctionTypeFlags<int_type>
+  withCoroutine(bool newValue = true) const {
+    return TargetExtendedFunctionTypeFlags<int_type>(
+        (Data & ~IsCoroutine) | (newValue ? IsCoroutine : 0));
+  }
+
   const TargetExtendedFunctionTypeFlags<int_type>
   withInvertedProtocols(InvertibleProtocolSet inverted) const {
     return TargetExtendedFunctionTypeFlags<int_type>(
@@ -1359,6 +1368,8 @@ class TargetExtendedFunctionTypeFlags {
     return bool(Data & HasSendingResult);
   }
 
+  bool isCoroutine() const { return bool(Data & IsCoroutine); }
+
   int_type getIntValue() const {
     return Data;
   }

include/swift/AST/ASTBridging.h

@@ -1473,14 +1473,16 @@ void BridgedDestructorDecl_setParsedBody(BridgedDestructorDecl decl,
 SWIFT_NAME("BridgedFuncDecl.createParsed(_:declContext:staticLoc:"
            "staticSpelling:funcKeywordLoc:"
            "name:nameLoc:genericParamList:parameterList:asyncSpecifierLoc:"
-           "throwsSpecifierLoc:thrownType:returnType:genericWhereClause:)")
+           "throwsSpecifierLoc:thrownType:yieldsKeywordLoc:yieldType:"
+           "returnType:genericWhereClause:)")
 BridgedFuncDecl BridgedFuncDecl_createParsed(
     BridgedASTContext cContext, BridgedDeclContext cDeclContext,
     swift::SourceLoc staticLoc, BridgedStaticSpelling cStaticSpelling,
     swift::SourceLoc funcKeywordLoc, swift::Identifier name,
     swift::SourceLoc nameLoc, BridgedNullableGenericParamList genericParamList,
     BridgedParameterList parameterList, swift::SourceLoc asyncLoc,
     swift::SourceLoc throwsLoc, BridgedNullableTypeRepr thrownType,
+    swift::SourceLoc yieldsLoc, BridgedNullableTypeRepr yieldType,
     BridgedNullableTypeRepr returnType,
     BridgedNullableTrailingWhereClause opaqueGenericWhereClause);
 
@@ -2647,11 +2649,12 @@ BridgedErrorTypeRepr BridgedErrorTypeRepr_create(BridgedASTContext cContext,
                                                  swift::SourceRange range);
 
 SWIFT_NAME("BridgedFunctionTypeRepr.createParsed(_:argsType:asyncLoc:throwsLoc:"
-           "thrownType:arrowLoc:resultType:)")
+           "thrownType:yieldsLoc:yieldType:arrowLoc:resultType:)")
 BridgedFunctionTypeRepr BridgedFunctionTypeRepr_createParsed(
     BridgedASTContext cContext, BridgedTypeRepr argsTy,
     swift::SourceLoc asyncLoc, swift::SourceLoc throwsLoc,
-    BridgedNullableTypeRepr thrownType, swift::SourceLoc arrowLoc,
+    BridgedNullableTypeRepr thrownType, swift::SourceLoc yieldsLoc,
+    BridgedNullableTypeRepr yieldType, swift::SourceLoc arrowLoc,
     BridgedTypeRepr resultType);
 
 SWIFT_NAME("BridgedUnqualifiedIdentTypeRepr.createParsed(_:name:nameLoc:"

include/swift/AST/ASTPrinter.h

@@ -105,6 +105,9 @@ enum class PrintStructureKind {
   FunctionParameterList,
   /// '@attribute ParamTy...' in parameter declarations.
   FunctionParameterType,
+  CoroutineYield,
+  /// 'yields Tys...`
+  CoroutineYieldsTypes,
 };
 
 /// ---------------------------------

include/swift/AST/ASTSynthesis.h

@@ -512,6 +512,34 @@ void synthesizeParameterTypes(SynthesisContext &SC,
   list.params.visit(CollectParamTypes{SC, types});
 }
 
+/// Synthesize a yield list.
+// TBD: Add helpers for inout yields, etc.
+template <class... Yields>
+struct YieldListSynthesizer {
+  VariadicSynthesizerStorage<Yields...> params;
+};
+template <class... Yields>
+constexpr YieldListSynthesizer<Yields...> _yields(Yields... ys) {
+  return {{ys...}};
+}
+
+struct CollectYieldTypes {
+  SynthesisContext &SC;
+  SmallVectorImpl<FunctionType::Yield> &Results;
+
+  template <class S>
+  void operator()(const S &s) const {
+    // Found by argument-dependent lookup.
+    Results.push_back(synthesizeYieldType(SC, s));
+  }
+};
+template <class... Yields>
+void synthesizeYieldTypes(SynthesisContext &SC,
+                          const YieldListSynthesizer<Yields...> &list,
+                          SmallVectorImpl<FunctionType::Yield> &types) {
+  list.params.visit(CollectYieldTypes{SC, types});
+}
+
 /// Synthesize function ExtInfo.
 template <class S> struct ThrowsSynthesizer { S sub; };
 template <class S> struct AsyncSynthesizer { S sub; };
@@ -557,25 +585,29 @@ ASTExtInfo synthesizeExtInfo(SynthesisContext &SC,
 }
 
 /// Synthesize a function type.
-template <class ExtInfoS, class ResultS, class ParamsS>
+template <class ExtInfoS, class ResultS, class ParamsS, class YieldsS>
 struct FunctionTypeSynthesizer {
   ExtInfoS extInfo;
   ResultS result;
   ParamsS parameters;
+  YieldsS yields;
 };
-template <class ExtInfoS, class ResultS, class ParamsS>
-FunctionTypeSynthesizer<ExtInfoS, ResultS, ParamsS>
-_function(ExtInfoS extInfo, ResultS result, ParamsS params) {
-  return {extInfo, result, params};
-}
-template <class ExtInfoS, class ResultS, class ParamsS>
-Type synthesizeType(SynthesisContext &SC,
-             const FunctionTypeSynthesizer<ExtInfoS, ResultS, ParamsS> &fn) {
+template <class ExtInfoS, class ResultS, class ParamsS, class YieldsS>
+FunctionTypeSynthesizer<ExtInfoS, ResultS, ParamsS, YieldsS>
+_function(ExtInfoS extInfo, ResultS result, ParamsS params, YieldsS yields) {
+  return {extInfo, result, params, yields};
+}
+template <class ExtInfoS, class ResultS, class ParamsS, class YieldsS>
+Type synthesizeType(
+    SynthesisContext &SC,
+    const FunctionTypeSynthesizer<ExtInfoS, ResultS, ParamsS, YieldsS> &fn) {
   SmallVector<FunctionType::Param, 4> paramTypes;
   synthesizeParameterTypes(SC, fn.parameters, paramTypes);
+  SmallVector<FunctionType::Yield, 1> yieldTypes;
+  synthesizeYieldTypes(SC, fn.yields, yieldTypes);
   auto extInfo = synthesizeExtInfo(SC, fn.extInfo);
   auto resultType = synthesizeType(SC, fn.result);
-  return FunctionType::get(paramTypes, resultType, extInfo);
+  return FunctionType::get(paramTypes, yieldTypes, resultType, extInfo);
 }
 
 /// Synthesize optionals.

include/swift/AST/AnyFunctionRef.h

@@ -300,25 +300,7 @@ class AnyFunctionRef {
 private:
   ArrayRef<AnyFunctionType::Yield>
   getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buffer,
-                      bool mapIntoContext) const {
-    assert(buffer.empty());
-    if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
-      if (auto *AD = dyn_cast<AccessorDecl>(AFD)) {
-        if (AD->isCoroutine()) {
-          auto valueTy = AD->getStorage()->getValueInterfaceType()
-                                         ->getReferenceStorageReferent();
-          if (mapIntoContext)
-            valueTy = AD->mapTypeIntoEnvironment(valueTy);
-          YieldTypeFlags flags(isYieldingMutableAccessor(AD->getAccessorKind())
-                                   ? ParamSpecifier::InOut
-                                   : ParamSpecifier::LegacyShared);
-          buffer.push_back(AnyFunctionType::Yield(valueTy, flags));
-          return buffer;
-        }
-      }
-    }
-    return {};
-  }
+                      bool mapIntoContext) const;
 };
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(pop)

include/swift/AST/Decl.h

@@ -7930,6 +7930,8 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// attribute.
   bool isTransparent() const;
 
+  bool isCoroutine() const;
+
   // Expose our import as member status
   ImportAsMemberStatus getImportAsMemberStatus() const {
     return ImportAsMemberStatus(Bits.AbstractFunctionDecl.IAMStatus);
@@ -8441,11 +8443,13 @@ class FuncDecl : public AbstractFunctionDecl {
   friend class SelfAccessKindRequest;
   friend class IsStaticRequest;
   friend class ResultTypeRequest;
+  friend class YieldsTypeRequest;
 
   SourceLoc StaticLoc;  // Location of the 'static' token or invalid.
   SourceLoc FuncLoc;    // Location of the 'func' token.
 
   TypeLoc FnRetType;
+  TypeLoc FnYieldType;
 
 protected:
   FuncDecl(DeclKind Kind,
@@ -8478,6 +8482,7 @@ class FuncDecl : public AbstractFunctionDecl {
   }
 
   void setResultInterfaceType(Type type);
+  void setYieldInterfaceType(Type type);
 
 private:
   static FuncDecl *createImpl(ASTContext &Context, SourceLoc StaticLoc,
@@ -8510,16 +8515,16 @@ class FuncDecl : public AbstractFunctionDecl {
   static FuncDecl *createDeserialized(ASTContext &Context,
                                       StaticSpellingKind StaticSpelling,
                                       DeclName Name, bool Async, bool Throws,
-                                      Type ThrownType,
+                                      Type ThrownType, Type YieldType,
                                       GenericParamList *GenericParams,
                                       Type FnRetType, DeclContext *Parent);
 
   static FuncDecl *create(ASTContext &Context, SourceLoc StaticLoc,
                           StaticSpellingKind StaticSpelling, SourceLoc FuncLoc,
                           DeclName Name, SourceLoc NameLoc, bool Async,
                           SourceLoc AsyncLoc, bool Throws, SourceLoc ThrowsLoc,
-                          TypeRepr *ThrownTyR,
-                          GenericParamList *GenericParams,
+                          TypeRepr *ThrownTyR, SourceLoc YieldsLoc,
+                          TypeRepr *YieldTyR, GenericParamList *GenericParams,
                           ParameterList *BodyParams, TypeRepr *ResultTyR,
                           DeclContext *Parent);
 
@@ -8576,18 +8581,24 @@ class FuncDecl : public AbstractFunctionDecl {
   SourceRange getSourceRange() const;
 
   TypeRepr *getResultTypeRepr() const { return FnRetType.getTypeRepr(); }
+  TypeRepr *getYieldTypeRepr() const { return FnYieldType.getTypeRepr(); }
 
   SourceRange getResultTypeSourceRange() const {
     return FnRetType.getSourceRange();
   }
 
-  /// Retrieve the result interface type of this function.
+  /// Retrieve the result interface type of this function
   Type getResultInterfaceType() 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;
 
+  std::optional<Type> getCachedYieldsInterfaceType() const;
+
   /// isUnaryOperator - Determine whether this is a unary operator
   /// implementation.  This check is a syntactic rather than type-based check,
   /// which looks at the number of parameters specified, in order to allow

include/swift/AST/DeclAttr.def

@@ -916,7 +916,12 @@ DECL_ATTR(warn, Warn,
   AllowMultipleAttributes | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove | ForbiddenInABIAttr,
   174)
 
-LAST_DECL_ATTR(Warn)
+SIMPLE_DECL_ATTR(yield_once, Coroutine,
+  OnFunc,
+  UserInaccessible | ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove | EquivalentInABIAttr,
+  175)
+
+LAST_DECL_ATTR(Coroutine)
 
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/DiagnosticsParse.def

@@ -331,6 +331,8 @@ ERROR(accessor_requires_coroutine_accessors,none,
 ERROR(accessor_requires_borrow_and_mutate_accessors,none,
       "%0 requires '-enable-experimental-feature BorrowAndMutateAccessors'",
       (StringRef))
+ERROR(yields_requires_coroutine_functions,none,
+      "`yields` requires '-enable-experimental-feature CoroutineFunctions'", ())
 
 // Import
 ERROR(decl_expected_module_name,none,
@@ -1724,6 +1726,11 @@ ERROR(expected_rparen_after_thrown_error_type,none,
 ERROR(rethrows_with_thrown_error,none,
       "'rethrows' cannot be combined with a specific thrown error type", ())
 
+ERROR(expected_yield_type,none,
+      "expected yield type after 'yields('", ())
+ERROR(expected_rparen_after_yield_type,none,
+      "expected ')' after yield type", ())
+
 ERROR(attr_isolated_expected_lparen,none,
       "expected '(' after '@isolated'; did you mean the 'isolated' modifier?",
       ())

include/swift/AST/Expr.h

@@ -5533,35 +5533,39 @@ class CoerceExpr final : public ExplicitCastExpr {
 class ArrowExpr : public Expr {
   SourceLoc AsyncLoc;
   SourceLoc ThrowsLoc;
+  SourceLoc YieldsLoc;
   SourceLoc ArrowLoc;
   Expr *Args;
+  Expr *Yield;
   Expr *Result;
   Expr *ThrownType;
 
 public:
   ArrowExpr(Expr *Args, SourceLoc AsyncLoc, SourceLoc ThrowsLoc,
-            Expr *ThrownType, SourceLoc ArrowLoc, Expr *Result)
-    : Expr(ExprKind::Arrow, /*implicit=*/false, Type()),
-      AsyncLoc(AsyncLoc), ThrowsLoc(ThrowsLoc), ArrowLoc(ArrowLoc), Args(Args),
-      Result(Result), ThrownType(ThrownType)
-  { }
+            Expr *ThrownType, SourceLoc YieldsLoc, Expr *Yield,
+            SourceLoc ArrowLoc, Expr *Result)
+      : Expr(ExprKind::Arrow, /*implicit=*/false, Type()), AsyncLoc(AsyncLoc),
+        ThrowsLoc(ThrowsLoc), YieldsLoc(YieldsLoc), ArrowLoc(ArrowLoc),
+        Args(Args), Yield(Yield), Result(Result), ThrownType(ThrownType) {}
 
   ArrowExpr(SourceLoc AsyncLoc, SourceLoc ThrowsLoc, Expr *ThrownType,
             SourceLoc ArrowLoc)
-    : Expr(ExprKind::Arrow, /*implicit=*/false, Type()),
-      AsyncLoc(AsyncLoc), ThrowsLoc(ThrowsLoc), ArrowLoc(ArrowLoc),
-      Args(nullptr), Result(nullptr), ThrownType(ThrownType)
-  { }
+      : Expr(ExprKind::Arrow, /*implicit=*/false, Type()), AsyncLoc(AsyncLoc),
+        ThrowsLoc(ThrowsLoc), ArrowLoc(ArrowLoc), Args(nullptr), Yield(nullptr),
+        Result(nullptr), ThrownType(ThrownType) {}
 
   Expr *getArgsExpr() const { return Args; }
   void setArgsExpr(Expr *E) { Args = E; }
+  Expr *getYieldExpr() const { return Yield; }
+  void setYieldExpr(Expr *E) { Yield = E; }
   Expr *getResultExpr() const { return Result; }
   void setResultExpr(Expr *E) { Result = E; }
   Expr *getThrownTypeExpr() const { return ThrownType; }
   void setThrownTypeExpr(Expr *E) { ThrownType = E; }
 
   SourceLoc getAsyncLoc() const { return AsyncLoc; }
   SourceLoc getThrowsLoc() const { return ThrowsLoc; }
+  SourceLoc getYieldsLoc() const { return YieldsLoc; }
   SourceLoc getArrowLoc() const { return ArrowLoc; }
   bool isFolded() const { return Args != nullptr && Result != nullptr; }