Basic motions towards moving yields into AnyFunctionType.

Still not intending to do this anytime soon.

Author: rjmccall

include/swift/AST/AnyFunctionRef.h

@@ -83,29 +83,14 @@ class AnyFunctionRef {
     return TheFunction.get<AbstractClosureExpr *>()->getResultType();
   }
 
-  struct YieldResult {
-    Type Ty;
-    VarDecl::Specifier Specifier;
-  };
+  ArrayRef<AnyFunctionType::Yield>
+  getYieldResults(SmallVectorImpl<AnyFunctionType::Yield> &buffer) const {
+    return getYieldResultsImpl(buffer, /*mapIntoContext*/ false);
+  }
 
-  ArrayRef<YieldResult>
-  getBodyYieldResults(SmallVectorImpl<YieldResult> &buffer) 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();
-          valueTy = AD->mapTypeIntoContext(valueTy);
-          auto specifier =
-            AD->getAccessorKind() == AccessorKind::Modify
-              ? VarDecl::Specifier::InOut
-              : VarDecl::Specifier::Shared;
-          buffer.push_back({valueTy, specifier});
-          return buffer;
-        }
-      }
-    }
-    return {};
+  ArrayRef<AnyFunctionType::Yield>
+  getBodyYieldResults(SmallVectorImpl<AnyFunctionType::Yield> &buffer) const {
+    return getYieldResultsImpl(buffer, /*mapIntoContext*/ true);
   }
 
   BraceStmt *getBody() const {
@@ -196,6 +181,28 @@ class AnyFunctionRef {
     }
     llvm_unreachable("unexpected AnyFunctionRef representation");
   }
+
+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();
+          if (mapIntoContext)
+            valueTy = AD->mapTypeIntoContext(valueTy);
+          YieldTypeFlags flags(AD->getAccessorKind() == AccessorKind::Modify
+                                 ? ValueOwnership::InOut
+                                 : ValueOwnership::Shared);
+          buffer.push_back(AnyFunctionType::Yield(valueTy, flags));
+          return buffer;
+        }
+      }
+    }
+    return {};
+  }
 };
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(pop)

include/swift/AST/Types.h

@@ -1745,6 +1745,70 @@ class ParameterTypeFlags {
   uint8_t toRaw() const { return value.toRaw(); }
 };
 
+class YieldTypeFlags {
+  enum YieldFlags : uint8_t {
+    None        = 0,
+    InOut       = 1 << 1,
+    Shared      = 1 << 2,
+    Owned       = 1 << 3,
+
+    NumBits = 3
+  };
+  OptionSet<YieldFlags> value;
+
+  static_assert(NumBits < 8 * sizeof(OptionSet<YieldFlags>), "overflowed");
+
+  YieldTypeFlags(OptionSet<YieldFlags, uint8_t> val) : value(val) {}
+
+public:
+  YieldTypeFlags() = default;
+  static YieldTypeFlags fromRaw(uint8_t raw) {
+    return YieldTypeFlags(OptionSet<YieldFlags>(raw));
+  }
+
+  YieldTypeFlags(ValueOwnership ownership)
+      : value((ownership == ValueOwnership::InOut ? InOut : 0) |
+              (ownership == ValueOwnership::Shared ? Shared : 0) |
+              (ownership == ValueOwnership::Owned ? Owned : 0)) {}
+
+  bool isInOut() const { return value.contains(InOut); }
+  bool isShared() const { return value.contains(Shared); }
+  bool isOwned() const { return value.contains(Owned); }
+
+  ValueOwnership getValueOwnership() const {
+    if (isInOut())
+      return ValueOwnership::InOut;
+    else if (isShared())
+      return ValueOwnership::Shared;
+    else if (isOwned())
+      return ValueOwnership::Owned;
+
+    return ValueOwnership::Default;
+  }
+
+  YieldTypeFlags withInOut(bool isInout) const {
+    return YieldTypeFlags(isInout ? value | InOut : value - InOut);
+  }
+  
+  YieldTypeFlags withShared(bool isShared) const {
+    return YieldTypeFlags(isShared ? value | Shared : value - Shared);
+  }
+
+  YieldTypeFlags withOwned(bool isOwned) const {
+    return YieldTypeFlags(isOwned ? value | Owned : value - Owned);
+  }
+
+  bool operator ==(const YieldTypeFlags &other) const {
+    return value.toRaw() == other.value.toRaw();
+  }
+
+  bool operator!=(const YieldTypeFlags &other) const {
+    return value.toRaw() != other.value.toRaw();
+  }
+
+  uint8_t toRaw() const { return value.toRaw(); }
+};
+
 /// ParenType - A paren type is a type that's been written in parentheses.
 class ParenType : public SugarType {
   friend class ASTContext;
@@ -2560,7 +2624,7 @@ class AnyFunctionType : public TypeBase {
 
 public:
   using Representation = FunctionTypeRepresentation;
-  
+
   class Param {
   public:
     explicit Param(const TupleTypeElt &tte);
@@ -2635,6 +2699,50 @@ class AnyFunctionType : public TypeBase {
   using CanParamArrayRef =
     ArrayRefView<Param,CanParam,CanParam::getFromParam,/*AccessOriginal*/true>;
 
+  class CanYield;
+  class Yield {
+    Type Ty;
+    YieldTypeFlags Flags;
+  public:
+    explicit Yield(Type type, YieldTypeFlags flags)
+      : Ty(type), Flags(flags) {}
+
+    Type getType() const { return Ty; }
+
+    YieldTypeFlags getFlags() const { return Flags; }
+    ValueOwnership getValueOwnership() const {
+      return getFlags().getValueOwnership();
+    }
+    bool isInOut() const { return getFlags().isInOut(); }
+
+    CanYield getCanonical() const;
+
+    Yield subst(SubstitutionMap subs, SubstOptions options = None) const {
+      return Yield(getType().subst(subs, options), getFlags());
+    }
+
+    bool operator==(const Yield &other) const {
+      return getType()->isEqual(other.getType()) &&
+             getFlags() == other.getFlags();
+    }
+    bool operator!=(const Yield &other) const {
+      return !operator==(other);
+    }
+  };
+
+  class CanYield : public Yield {
+  public:
+    explicit CanYield(CanType type, YieldTypeFlags flags)
+      : Yield(type, flags) {}
+
+    CanType getType() const { return CanType(Yield::getType()); }
+
+    CanYield subst(SubstitutionMap subs, SubstOptions options = None) const {
+      return CanYield(getType().subst(subs, options)->getCanonicalType(),
+                      getFlags());
+    }
+  };
+
   /// \brief A class which abstracts out some details necessary for
   /// making a call.
   class ExtInfo {
@@ -2896,6 +3004,10 @@ BEGIN_CAN_TYPE_WRAPPER(AnyFunctionType, Type)
   }
 END_CAN_TYPE_WRAPPER(AnyFunctionType, Type)
 
+inline AnyFunctionType::CanYield AnyFunctionType::Yield::getCanonical() const {
+  return CanYield(getType()->getCanonicalType(), getFlags());
+}
+
 /// FunctionType - A monomorphic function type, specified with an arrow.
 ///
 /// For example:

lib/Sema/TypeCheckStmt.cpp

@@ -482,7 +482,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       return YS;
     }
 
-    SmallVector<AnyFunctionRef::YieldResult, 4> buffer;
+    SmallVector<AnyFunctionType::Yield, 4> buffer;
     auto yieldResults = TheFunc->getBodyYieldResults(buffer);
 
     auto yieldExprs = YS->getMutableYields();
@@ -493,15 +493,15 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
     }
 
     for (auto i : indices(yieldExprs)) {
-      Type yieldType = yieldResults[i].Ty;
+      Type yieldType = yieldResults[i].getType();
       auto exprToCheck = yieldExprs[i];
 
       InOutExpr *inout = nullptr;
 
       // Classify whether we're yielding by reference or by value.
       ContextualTypePurpose contextTypePurpose;
       Type contextType = yieldType;
-      if (yieldResults[i].Specifier == VarDecl::Specifier::InOut) {
+      if (yieldResults[i].isInOut()) {
         contextTypePurpose = CTP_YieldByReference;
         contextType = LValueType::get(contextType);