Merge pull request #3235 from slavapestov/fix-capturing-dynamic-self

Fix capturing dynamic Self

Author: slavapestov

include/swift/AST/CaptureInfo.h

@@ -55,10 +55,20 @@ class CapturedValue {
 
   CapturedValue(ValueDecl *D, unsigned Flags) : Value(D, Flags) {}
 
+  static CapturedValue getDynamicSelfMetadata() {
+    return CapturedValue(nullptr, 0);
+  }
+
   bool isDirect() const { return Value.getInt() & IsDirect; }
   bool isNoEscape() const { return Value.getInt() & IsNoEscape; }
 
-  ValueDecl *getDecl() const { return Value.getPointer(); }
+  bool isDynamicSelfMetadata() const { return !Value.getPointer(); }
+
+  ValueDecl *getDecl() const {
+    assert(Value.getPointer() && "dynamic Self metadata capture does not "
+           "have a value");
+    return Value.getPointer();
+  }
 
   unsigned getFlags() const { return Value.getInt(); }
 
@@ -106,19 +116,26 @@ template <> struct DenseMapInfo<swift::CapturedValue> {
 
 namespace swift {
 
+class DynamicSelfType;
+
 /// \brief Stores information about captured variables.
 class CaptureInfo {
   const CapturedValue *Captures;
+  DynamicSelfType *DynamicSelf;
   unsigned Count = 0;
   bool GenericParamCaptures : 1;
   bool Computed : 1;
 
 public:
   CaptureInfo()
-    : Captures(nullptr), Count(0), GenericParamCaptures(0), Computed(0) { }
+    : Captures(nullptr), DynamicSelf(nullptr), Count(0),
+      GenericParamCaptures(0), Computed(0) { }
 
   bool hasBeenComputed() { return Computed; }
-  bool empty() { return Count == 0; }
+
+  bool isTrivial() {
+    return Count == 0 && !GenericParamCaptures && !DynamicSelf;
+  }
 
   ArrayRef<CapturedValue> getCaptures() const {
     return llvm::makeArrayRef(Captures, Count);
@@ -139,14 +156,28 @@ class CaptureInfo {
   bool hasLocalCaptures() const;
 
   /// \returns true if the function captures any generic type parameters.
-  bool hasGenericParamCaptures() {
+  bool hasGenericParamCaptures() const {
     return GenericParamCaptures;
   }
 
   void setGenericParamCaptures(bool genericParamCaptures) {
     GenericParamCaptures = genericParamCaptures;
   }
 
+  /// \returns true if the function captures the dynamic Self type.
+  bool hasDynamicSelfCapture() const {
+    return DynamicSelf != nullptr;
+  }
+
+  /// \returns the captured dynamic Self type, if any.
+  DynamicSelfType *getDynamicSelfType() const {
+    return DynamicSelf;
+  }
+
+  void setDynamicSelfType(DynamicSelfType *dynamicSelf) {
+    DynamicSelf = dynamicSelf;
+  }
+
   void dump() const;
   void print(raw_ostream &OS) const;
 };

include/swift/AST/DiagnosticsSema.def

@@ -802,7 +802,8 @@ ERROR(c_function_pointer_from_generic_function,none,
 ERROR(c_function_pointer_from_function_with_context,none,
       "a C function pointer cannot be formed from a "
       "%select{local function|closure}0 that captures "
-      "%select{context|generic parameters}1", (bool, bool))
+      "%select{context|generic parameters|dynamic Self type}1",
+      (bool, unsigned))
 NOTE(c_function_pointer_captures_here,none,
      "%0 captured here", (Identifier))
 

lib/AST/ASTDumper.cpp

@@ -737,7 +737,7 @@ namespace {
 
     void printCommonAFD(AbstractFunctionDecl *D, const char *Type) {
       printCommon(D, Type, FuncColor);
-      if (!D->getCaptureInfo().empty()) {
+      if (!D->getCaptureInfo().isTrivial()) {
         OS << " ";
         D->getCaptureInfo().print(OS);
       }
@@ -2014,7 +2014,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
   llvm::raw_ostream &printClosure(AbstractClosureExpr *E, char const *name) {
     printCommon(E, name);
     OS << " discriminator=" << E->getDiscriminator();
-    if (!E->getCaptureInfo().empty()) {
+    if (!E->getCaptureInfo().isTrivial()) {
       OS << " ";
       E->getCaptureInfo().print(OS);
     }

lib/AST/CaptureInfo.cpp

@@ -46,6 +46,12 @@ void CaptureInfo::dump() const {
 
 void CaptureInfo::print(raw_ostream &OS) const {
   OS << "captures=(";
+
+  if (hasGenericParamCaptures())
+    OS << "<generic> ";
+  if (hasDynamicSelfCapture())
+    OS << "<dynamic_self> ";
+
   bool isFirst = true;
 
   for (auto capture : getCaptures()) {

lib/SIL/SILArgument.cpp

@@ -26,12 +26,6 @@ using namespace swift;
 SILArgument::SILArgument(SILBasicBlock *ParentBB, SILType Ty,
                          const ValueDecl *D)
   : ValueBase(ValueKind::SILArgument, Ty), ParentBB(ParentBB), Decl(D) {
-  // Function arguments need to have a decl.
-  assert(
-    !ParentBB->getParent()->isBare() &&
-    ParentBB->getParent()->size() == 1
-          ? D != nullptr
-          : true);
   ParentBB->insertArgument(ParentBB->bbarg_end(), this);
 }
 

lib/SIL/SILFunctionType.cpp

@@ -735,6 +735,18 @@ static CanSILFunctionType getSILFunctionType(SILModule &M,
     auto loweredCaptures = Types.getLoweredLocalCaptures(*function);
 
     for (auto capture : loweredCaptures.getCaptures()) {
+      if (capture.isDynamicSelfMetadata()) {
+        ParameterConvention convention = ParameterConvention::Direct_Unowned;
+        auto selfMetatype = MetatypeType::get(
+            loweredCaptures.getDynamicSelfType()->getSelfType(),
+            MetatypeRepresentation::Thick)
+                ->getCanonicalType();
+        SILParameterInfo param(selfMetatype, convention);
+        inputs.push_back(param);
+
+        continue;
+      }
+
       auto *VD = capture.getDecl();
       auto type = VD->getType()->getCanonicalType();
 

lib/SIL/SILVerifier.cpp

@@ -1326,6 +1326,23 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     if (auto dynamicSelf = dyn_cast<DynamicSelfType>(formalType))
       formalType = CanType(dynamicSelf->getSelfType());
 
+    // Optional of dynamic self has the same lowering as its contained type.
+    OptionalTypeKind loweredOptionalKind;
+    OptionalTypeKind formalOptionalKind;
+
+    CanType loweredObjectType = loweredType.getSwiftRValueType()
+        .getAnyOptionalObjectType(loweredOptionalKind);
+    CanType formalObjectType = formalType
+        .getAnyOptionalObjectType(formalOptionalKind);
+
+    if (loweredOptionalKind != OTK_None) {
+      if (auto dynamicSelf = dyn_cast<DynamicSelfType>(formalObjectType)) {
+        formalObjectType = dynamicSelf->getSelfType()->getCanonicalType();
+      }
+      return ((loweredOptionalKind == formalOptionalKind) &&
+              loweredObjectType == formalObjectType);
+    }
+
     // Metatypes preserve their instance type through lowering.
     if (auto loweredMT = loweredType.getAs<MetatypeType>()) {
       if (auto formalMT = dyn_cast<MetatypeType>(formalType)) {

lib/SIL/TypeLowering.cpp

@@ -2074,13 +2074,14 @@ getAnyFunctionRefFromCapture(CapturedValue capture) {
 CaptureInfo
 TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
   // First, bail out if there are no local captures at all.
-  if (!fn.getCaptureInfo().hasLocalCaptures()) {
+  if (!fn.getCaptureInfo().hasLocalCaptures() &&
+      !fn.getCaptureInfo().hasDynamicSelfCapture()) {
     CaptureInfo info;
     info.setGenericParamCaptures(
         fn.getCaptureInfo().hasGenericParamCaptures());
     return info;
   };
-  
+
   // See if we've cached the lowered capture list for this function.
   auto found = LoweredCaptures.find(fn);
   if (found != LoweredCaptures.end())
@@ -2089,7 +2090,13 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
   // Recursively collect transitive captures from captured local functions.
   llvm::DenseSet<AnyFunctionRef> visitedFunctions;
   llvm::SetVector<CapturedValue> captures;
+
+  // If there is a capture of 'self' with dynamic 'Self' type, it goes last so
+  // that IRGen can pass dynamic 'Self' metadata.
+  Optional<CapturedValue> selfCapture;
+
   bool capturesGenericParams = false;
+  DynamicSelfType *capturesDynamicSelf = nullptr;
 
   std::function<void (AnyFunctionRef)> collectFunctionCaptures
   = [&](AnyFunctionRef curFn) {
@@ -2098,6 +2105,8 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
 
     if (curFn.getCaptureInfo().hasGenericParamCaptures())
       capturesGenericParams = true;
+    if (curFn.getCaptureInfo().hasDynamicSelfCapture())
+      capturesDynamicSelf = curFn.getCaptureInfo().getDynamicSelfType();
 
     SmallVector<CapturedValue, 4> localCaptures;
     curFn.getCaptureInfo().getLocalCaptures(localCaptures);
@@ -2108,7 +2117,7 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
         collectFunctionCaptures(*capturedFn);
         continue;
       }
-      
+
       // If the capture is of a computed property, grab the transitive captures
       // of its accessors.
       if (auto capturedVar = dyn_cast<VarDecl>(capture.getDecl())) {
@@ -2127,21 +2136,38 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
           // Directly capture storage if we're supposed to.
           if (capture.isDirect())
             goto capture_value;
-            
+
           // Otherwise, transitively capture the accessors.
           SWIFT_FALLTHROUGH;
-          
+
         case VarDecl::Computed: {
           collectFunctionCaptures(capturedVar->getGetter());
           if (auto setter = capturedVar->getSetter())
             collectFunctionCaptures(setter);
           continue;
         }
 
-        case VarDecl::Stored:
+        case VarDecl::Stored: {
           // We can always capture the storage in these cases.
+          Type captureType;
+          if (auto *selfType = capturedVar->getType()->getAs<DynamicSelfType>()) {
+            captureType = selfType->getSelfType();
+            if (auto *metatypeType = captureType->getAs<MetatypeType>())
+              captureType = metatypeType->getInstanceType();
+
+            // We're capturing a 'self' value with dynamic 'Self' type;
+            // handle it specially.
+            if (!selfCapture &&
+                captureType->getClassOrBoundGenericClass()) {
+              selfCapture = capture;
+              continue;
+            }
+          }
+
+          // Otherwise just fall through.
           goto capture_value;
         }
+        }
       }
 
     capture_value:
@@ -2150,12 +2176,23 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
     }
   };
   collectFunctionCaptures(fn);
-  
+
+  // If we captured the dynamic 'Self' type and we have a 'self' value also,
+  // add it as the final capture. Otherwise, add a fake hidden capture for
+  // the dynamic 'Self' metatype.
+  if (selfCapture.hasValue()) {
+    captures.insert(*selfCapture);
+  } else if (capturesDynamicSelf) {
+    selfCapture = CapturedValue::getDynamicSelfMetadata();
+    captures.insert(*selfCapture);
+  }
+
   // Cache the uniqued set of transitive captures.
   auto inserted = LoweredCaptures.insert({fn, CaptureInfo()});
   assert(inserted.second && "already in map?!");
   auto &cachedCaptures = inserted.first->second;
   cachedCaptures.setGenericParamCaptures(capturesGenericParams);
+  cachedCaptures.setDynamicSelfType(capturesDynamicSelf);
   cachedCaptures.setCaptures(Context.AllocateCopy(captures));
 
   return cachedCaptures;

lib/SILGen/SILGenFunction.cpp

@@ -232,6 +232,28 @@ void SILGenFunction::emitCaptures(SILLocation loc,
     canGuarantee = false;
 
   for (auto capture : captureInfo.getCaptures()) {
+    if (capture.isDynamicSelfMetadata()) {
+      // The parameter type is the static Self type, but the value we
+      // want to pass is the dynamic Self type, so upcast it.
+      auto dynamicSelfMetatype = MetatypeType::get(
+          captureInfo.getDynamicSelfType(),
+          MetatypeRepresentation::Thick)
+              ->getCanonicalType();
+      auto staticSelfMetatype = MetatypeType::get(
+          captureInfo.getDynamicSelfType()->getSelfType(),
+          MetatypeRepresentation::Thick)
+              ->getCanonicalType();
+      SILType dynamicSILType = SILType::getPrimitiveObjectType(
+          dynamicSelfMetatype);
+      SILType staticSILType = SILType::getPrimitiveObjectType(
+          staticSelfMetatype);
+
+      SILValue value = B.createMetatype(loc, dynamicSILType);
+      value = B.createUpcast(loc, value, staticSILType);
+      capturedArgs.push_back(ManagedValue::forUnmanaged(value));
+      continue;
+    }
+
     auto *vd = capture.getDecl();
 
     switch (SGM.Types.getDeclCaptureKind(capture)) {

lib/SILGen/SILGenProlog.cpp

@@ -357,6 +357,7 @@ void SILGenFunction::bindParametersForForwarding(const ParameterList *params,
 
 static void emitCaptureArguments(SILGenFunction &gen, CapturedValue capture,
                                  unsigned ArgNo) {
+
   auto *VD = capture.getDecl();
   auto type = VD->getType();
   SILLocation Loc(VD);
@@ -428,8 +429,21 @@ void SILGenFunction::emitProlog(AnyFunctionRef TheClosure,
   // Emit the capture argument variables. These are placed last because they
   // become the first curry level of the SIL function.
   auto captureInfo = SGM.Types.getLoweredLocalCaptures(TheClosure);
-  for (auto capture : captureInfo.getCaptures())
+  for (auto capture : captureInfo.getCaptures()) {
+    if (capture.isDynamicSelfMetadata()) {
+      auto selfMetatype = MetatypeType::get(
+          captureInfo.getDynamicSelfType()->getSelfType(),
+          MetatypeRepresentation::Thick)
+              ->getCanonicalType();
+      SILType ty = SILType::getPrimitiveObjectType(selfMetatype);
+      SILValue val = new (SGM.M) SILArgument(F.begin(), ty);
+      (void) val;
+
+      return;
+    }
+
     emitCaptureArguments(*this, capture, ++ArgNo);
+  }
 }
 
 static void emitIndirectResultParameters(SILGenFunction &gen, Type resultType,