Add basic boilerplate for AST coroutines and yields

Author: asl

include/swift/AST/Decl.h

@@ -7866,6 +7866,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);

include/swift/AST/DeclAttr.def

@@ -899,7 +899,12 @@ DECL_ATTR(specialized, Specialized,
   AllowMultipleAttributes | LongAttribute | UserInaccessible | ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove | ForbiddenInABIAttr,
   172)
 
-LAST_DECL_ATTR(Specialized)
+SIMPLE_DECL_ATTR(yield_once, Coroutine,
+  OnFunc,
+  UserInaccessible | ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove | EquivalentInABIAttr,
+  173)
+
+LAST_DECL_ATTR(Coroutine)
 
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/ExtInfo.h

@@ -323,7 +323,7 @@ enum class SILFunctionTypeRepresentation : uint8_t {
   CFunctionPointer = uint8_t(FunctionTypeRepresentation::CFunctionPointer),
 
   /// The value of the greatest AST function representation.
-  LastAST = CFunctionPointer,
+  LastAST = uint8_t(FunctionTypeRepresentation::Last),
 
   /// The value of the least SIL-only function representation.
   FirstSIL = 8,
@@ -523,8 +523,8 @@ class ASTExtInfoBuilder {
   // If bits are added or removed, then TypeBase::NumAFTExtInfoBits
   // and NumMaskBits must be updated, and they must match.
   //
-  //   |representation|noEscape|concurrent|async|throws|isolation|differentiability| SendingResult |
-  //   |    0 .. 3    |    4   |    5     |  6  |   7  | 8 .. 10 |     11 .. 13    |         14    |
+  //   |representation|noEscape|concurrent|async|throws|isolation|differentiability| SendingResult | coroutine |
+  //   |    0 .. 3    |    4   |    5     |  6  |   7  | 8 .. 10 |     11 .. 13    |         14    |    15     |
   //
   enum : unsigned {
     RepresentationMask = 0xF << 0,
@@ -537,7 +537,8 @@ class ASTExtInfoBuilder {
     DifferentiabilityMaskOffset = 11,
     DifferentiabilityMask = 0x7 << DifferentiabilityMaskOffset,
     SendingResultMask = 1 << 14,
-    NumMaskBits = 15
+    CoroutineMask = 1 << 15,
+    NumMaskBits = 16
   };
 
   static_assert(FunctionTypeIsolation::Mask == 0x7, "update mask manually");
@@ -616,6 +617,8 @@ class ASTExtInfoBuilder {
 
   constexpr bool hasSendingResult() const { return bits & SendingResultMask; }
 
+  constexpr bool isCoroutine() const { return bits & CoroutineMask; }
+
   constexpr DifferentiabilityKind getDifferentiabilityKind() const {
     return DifferentiabilityKind((bits & DifferentiabilityMask) >>
                                  DifferentiabilityMaskOffset);
@@ -732,6 +735,13 @@ class ASTExtInfoBuilder {
         clangTypeInfo, globalActor, thrownError, lifetimeDependencies);
   }
 
+  [[nodiscard]]
+  ASTExtInfoBuilder withCoroutine(bool coroutine = true) const {
+    return ASTExtInfoBuilder(
+        coroutine ? (bits | CoroutineMask) : (bits & ~CoroutineMask),
+        clangTypeInfo, globalActor, thrownError, lifetimeDependencies);
+  }
+
   [[nodiscard]]
   ASTExtInfoBuilder
   withDifferentiabilityKind(DifferentiabilityKind differentiability) const {
@@ -854,6 +864,8 @@ class ASTExtInfo {
 
   constexpr bool isThrowing() const { return builder.isThrowing(); }
 
+  constexpr bool isCoroutine() const { return builder.isCoroutine(); }
+
   constexpr bool hasSendingResult() const { return builder.hasSendingResult(); }
 
   constexpr DifferentiabilityKind getDifferentiabilityKind() const {
@@ -917,6 +929,14 @@ class ASTExtInfo {
     return builder.withThrows(true, Type()).build();
   }
 
+  /// Helper method for changing only the coroutine field.
+  ///
+  /// Prefer using \c ASTExtInfoBuilder::withCoroutine for chaining.
+  [[nodiscard]]
+  ASTExtInfo withCoroutine(bool coroutine = true) const {
+    return builder.withCoroutine(coroutine).build();
+  }
+
   /// Helper method for changing only the async field.
   ///
   /// Prefer using \c ASTExtInfoBuilder::withAsync for chaining.

include/swift/AST/TypeDifferenceVisitor.h

@@ -172,6 +172,10 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
                                type1->getElements(), type2->getElements());
   }
 
+  bool visitYieldResultType(CanYieldResultType type1, CanYieldResultType type2) {
+    return asImpl().visit(type1.getResultType(), type2.getResultType());
+  }
+
   bool visitComponent(CanType type1, CanType type2,
                       const TupleTypeElt &elt1, const TupleTypeElt &elt2) {
     if (elt1.getName() != elt2.getName())

include/swift/AST/TypeMatcher.h

@@ -153,6 +153,17 @@ class TypeMatcher {
       return mismatch(firstTuple.getPointer(), secondType, sugaredFirstType);
     }
 
+    bool visitYieldResultType(CanYieldResultType firstType, Type secondType,
+                              Type sugaredFirstType) {
+      if (auto secondYieldType = secondType->getAs<YieldResultType>())
+        if (!this->visit(firstType.getResultType(),
+                         secondYieldType->getResultType(),
+                         sugaredFirstType->getAs<YieldResultType>()->getResultType()))
+          return false;
+
+      return mismatch(firstType.getPointer(), secondType, sugaredFirstType);
+    }
+
     bool visitSILPackType(CanSILPackType firstPack, Type secondType,
                           Type sugaredFirstType) {
       if (auto secondPack = secondType->getAs<SILPackType>()) {

include/swift/AST/TypeNodes.def

@@ -205,6 +205,7 @@ TYPE(InOut, Type)
 TYPE(Pack, Type)
 TYPE(PackExpansion, Type)
 TYPE(PackElement, Type)
+TYPE(YieldResult, Type)
 UNCHECKED_TYPE(TypeVariable, Type)
 UNCHECKED_TYPE(ErrorUnion, Type)
 TYPE(Integer, Type)

include/swift/AST/TypeTransform.h

@@ -1040,6 +1040,16 @@ case TypeKind::Id:
         t : InOutType::get(objectTy);
     }
 
+    case TypeKind::YieldResult: {
+      auto yield = cast<YieldResultType>(base);
+      auto objectTy = doIt(yield->getResultType(), TypePosition::Invariant);
+      if (!objectTy || objectTy->hasError())
+        return objectTy;
+
+      return objectTy.getPointer() == yield->getResultType().getPointer() ?
+        t : YieldResultType::get(objectTy, yield->isInOut());
+    }
+
     case TypeKind::Existential: {
       auto *existential = cast<ExistentialType>(base);
       auto constraint = doIt(existential->getConstraintType(), pos);

include/swift/AST/Types.h

@@ -407,7 +407,7 @@ class alignas(1 << TypeAlignInBits) TypeBase
   }
 
 protected:
-  enum { NumAFTExtInfoBits = 15 };
+  enum { NumAFTExtInfoBits = 16 };
   enum { NumSILExtInfoBits = 14 };
 
   // clang-format off
@@ -437,15 +437,15 @@ class alignas(1 << TypeAlignInBits) TypeBase
     HasCachedType : 1
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(AnyFunctionType, TypeBase, NumAFTExtInfoBits+1+1+1+1+16,
+  SWIFT_INLINE_BITFIELD_FULL(AnyFunctionType, TypeBase, NumAFTExtInfoBits+1+1+1+1+14,
     /// Extra information which affects how the function is called, like
     /// regparm and the calling convention.
     ExtInfoBits : NumAFTExtInfoBits,
     HasExtInfo : 1,
     HasClangTypeInfo : 1,
     HasThrownError : 1,
     HasLifetimeDependencies : 1,
-    NumParams : 15
+    NumParams : 14
   );
 
   SWIFT_INLINE_BITFIELD_FULL(ArchetypeType, TypeBase, 1+1+16,
@@ -1702,6 +1702,33 @@ class UnresolvedType : public TypeBase {
 };
 DEFINE_EMPTY_CAN_TYPE_WRAPPER(UnresolvedType, Type)
 
+class YieldResultType : public TypeBase {
+  Type ResultType;
+  bool InOut = false;
+
+  YieldResultType(Type objectTy, bool InOut, const ASTContext *canonicalContext,
+                  RecursiveTypeProperties properties)
+  : TypeBase(TypeKind::YieldResult, canonicalContext, properties),
+    ResultType(objectTy), InOut(InOut) {}
+
+public:
+  static YieldResultType *get(Type originalType, bool InOut);
+
+  Type getResultType() const { return ResultType; }
+  bool isInOut() const { return InOut; }
+
+  // Implement isa/cast/dyncast/etc.
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::YieldResult;
+  }
+};
+
+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));
+}
+END_CAN_TYPE_WRAPPER(YieldResultType, Type)
 
 /// BuiltinType - An abstract class for all the builtin types.
 class BuiltinType : public TypeBase {
@@ -3886,6 +3913,8 @@ class AnyFunctionType : public TypeBase {
 
   bool isThrowing() const { return getExtInfo().isThrowing(); }
 
+  bool isCoroutine() const { return getExtInfo().isCoroutine(); }
+
   bool hasSendingResult() const { return getExtInfo().hasSendingResult(); }
 
   bool hasEffect(EffectKind kind) const;

lib/AST/ASTContext.cpp

@@ -555,6 +555,8 @@ struct ASTContext::Implementation {
     llvm::DenseMap<uintptr_t, ReferenceStorageType*> ReferenceStorageTypes;
     llvm::DenseMap<Type, LValueType*> LValueTypes;
     llvm::DenseMap<Type, InOutType*> InOutTypes;
+    llvm::DenseMap<llvm::PointerIntPair<TypeBase*, 1, bool>,
+                   YieldResultType*> YieldResultTypes;
     llvm::DenseMap<std::pair<Type, void*>, DependentMemberType *>
       DependentMemberTypes;
     llvm::FoldingSet<ErrorUnionType> ErrorUnionTypes;
@@ -3336,6 +3338,7 @@ size_t ASTContext::Implementation::Arena::getTotalMemory() const {
     llvm::capacity_in_bytes(ReferenceStorageTypes) +
     llvm::capacity_in_bytes(LValueTypes) +
     llvm::capacity_in_bytes(InOutTypes) +
+    llvm::capacity_in_bytes(YieldResultTypes) +
     llvm::capacity_in_bytes(DependentMemberTypes) +
     llvm::capacity_in_bytes(EnumTypes) +
     llvm::capacity_in_bytes(StructTypes) +
@@ -3374,6 +3377,7 @@ void ASTContext::Implementation::Arena::dump(llvm::raw_ostream &os) const {
     SIZE_AND_BYTES(ReferenceStorageTypes);
     SIZE_AND_BYTES(LValueTypes);
     SIZE_AND_BYTES(InOutTypes);
+    SIZE_AND_BYTES(YieldResultTypes);
     SIZE_AND_BYTES(DependentMemberTypes);
     SIZE(ErrorUnionTypes);
     SIZE_AND_BYTES(PlaceholderTypes);
@@ -5683,6 +5687,28 @@ InOutType *InOutType::get(Type objectTy) {
                                           properties);
 }
 
+YieldResultType *YieldResultType::get(Type objectTy, bool InOut) {
+  auto properties = objectTy->getRecursiveProperties();
+  if (InOut) {
+    assert(!objectTy->is<LValueType>() && !objectTy->is<InOutType>() &&
+           "cannot have 'inout' or @lvalue wrapped inside an 'inout yield'");
+    properties &= ~RecursiveTypeProperties::IsLValue;
+  }
+
+  auto arena = getArena(properties);
+
+  auto &C = objectTy->getASTContext();
+  auto pair = llvm::PointerIntPair<TypeBase*, 1, bool>(objectTy.getPointer(),
+                                                       InOut);
+  auto &entry = C.getImpl().getArena(arena).YieldResultTypes[pair];
+  if (entry)
+    return entry;
+
+  const ASTContext *canonicalContext = objectTy->isCanonical() ? &C : nullptr;
+  return entry = new (C, arena) YieldResultType(objectTy, InOut, canonicalContext,
+                                                properties);
+}
+
 DependentMemberType *DependentMemberType::get(Type base, Identifier name) {
   auto properties = base->getRecursiveProperties();
   properties |= RecursiveTypeProperties::HasDependentMember;

lib/AST/ASTDumper.cpp

@@ -4967,6 +4967,7 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, Label>,
   TRIVIAL_ATTR_PRINTER(CompilerInitialized, compiler_initialized)
   TRIVIAL_ATTR_PRINTER(Consuming, consuming)
   TRIVIAL_ATTR_PRINTER(Convenience, convenience)
+  TRIVIAL_ATTR_PRINTER(Coroutine, coroutine)
   TRIVIAL_ATTR_PRINTER(DiscardableResult, discardable_result)
   TRIVIAL_ATTR_PRINTER(DisfavoredOverload, disfavored_overload)
   TRIVIAL_ATTR_PRINTER(DistributedActor, distributed_actor)
@@ -6065,6 +6066,13 @@ namespace {
       printFoot();
     }
 
+    void visitYieldResultType(YieldResultType *T, Label label) {
+      printCommon("yield", label);
+      printFlag(T->isInOut(), "inout");
+      printRec(T->getResultType(), Label::always("type"));
+      printFoot();
+    }
+
     TRIVIAL_TYPE_PRINTER(Unresolved, unresolved)
 
     void visitPlaceholderType(PlaceholderType *T, Label label) {
@@ -6485,6 +6493,7 @@ namespace {
         printFlag(T->isAsync(), "async");
         printFlag(T->isThrowing(), "throws");
         printFlag(T->hasSendingResult(), "sending_result");
+        printFlag(T->isCoroutine(), "@yield_once");
         if (T->isDifferentiable()) {
           switch (T->getDifferentiabilityKind()) {
           default: