SIL: Add support for capturing the dynamic Self metatype

If a closure captures the dynamic 'Self' type, but no value of type 'Self'
(for example, it is possible to have a weak capture of 'self'; if the weak
reference becomes nil, there's no way for the closure to get the dynamic
'Self' type from the value).

In this case, add a hidden argument of type $Self.Type, and pass in the
Self metatype.

Fixes <https://bugs.swift.org/browse/SR-1558> / <rdar://problem/22299905>.

Author: slavapestov

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

@@ -2082,8 +2082,6 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
     return info;
   };
 
-  assert(fn.getAsDeclContext()->getParent()->isLocalContext());
-
   // See if we've cached the lowered capture list for this function.
   auto found = LoweredCaptures.find(fn);
   if (found != LoweredCaptures.end())
@@ -2092,8 +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;
+  DynamicSelfType *capturesDynamicSelf = nullptr;
 
   std::function<void (AnyFunctionRef)> collectFunctionCaptures
   = [&](AnyFunctionRef curFn) {
@@ -2144,10 +2147,27 @@ TypeConverter::getLoweredLocalCaptures(AnyFunctionRef fn) {
           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:
@@ -2156,7 +2176,17 @@ 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?!");

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,

test/Interpreter/dynamic_self.swift

@@ -94,5 +94,44 @@ print("C() as class existential")
 // CHECK-NEXT: Destroying C
 callDynamicSelfClassExistential(C())
 
+print("-------------------------------")
+
+class Z {
+  let name: String
+
+  init(name: String) {
+    self.name = name
+  }
+
+  func testCaptures(x: Int) -> Self {
+    let fn1 = {
+      print("First: \(self.name)")
+    }
+    fn1()
+
+    let fn2 = { [weak self] in
+      if let strongSelf = self {
+        print("Second: \(strongSelf.name)")
+      }
+    }
+    fn2()
+
+    let fn3 = {
+      print("Third: \(self.name)")
+      print("Third: \(x)")
+    }
+    fn3()
+
+    return self
+  }
+
+}
+
+// CHECK: First: Leeloo
+// CHECK-NEXT: Second: Leeloo
+// CHECK-NEXT: Third: Leeloo
+// CHECK-NEXT: Third: 42
+Z(name: "Leeloo").testCaptures(x: 42)
+
 // CHECK-NEXT: Done
 print("Done")

test/SILGen/dynamic_self.swift

@@ -152,6 +152,51 @@ func testOptionalResult(v : OptionalResultInheritor) {
 // CHECK-NEXT: [[T4:%.*]] = unchecked_ref_cast [[T1]] : $OptionalResult to $OptionalResultInheritor
 // CHECK-NEXT: enum $Optional<OptionalResultInheritor>, #Optional.some!enumelt.1, [[T4]]
 
+class Z {
+
+  // CHECK-LABEL: sil hidden @_TFC12dynamic_self1Z23testDynamicSelfCapturesfT1xSi_DS0_
+  func testDynamicSelfCaptures(x: Int) -> Self {
+
+    // Single capture of 'self' type
+
+    // CHECK:      [[FN:%.*]] = function_ref @_TFFC12dynamic_self1Z23testDynamicSelfCapturesFT1xSi_DS0_U_FT_T_ : $@convention(thin) (@owned Z) -> ()
+    // CHECK-NEXT: strong_retain %1 : $Z
+    // CHECK-NEXT: partial_apply [[FN]](%1)
+    let fn1 = { _ = self }
+    fn1()
+
+    // Capturing 'self', but it's not the last capture. Make sure it ends
+    // up at the end of the list anyway
+
+    // CHECK:      [[FN:%.*]] = function_ref @_TFFC12dynamic_self1Z23testDynamicSelfCapturesFT1xSi_DS0_U0_FT_T_ : $@convention(thin) (Int, @owned Z) -> ()
+    // CHECK-NEXT: strong_retain %1 : $Z
+    // CHECK-NEXT: partial_apply [[FN]](%0, %1)
+    let fn2 = {
+      _ = self
+      _ = x
+    }
+    fn2()
+
+    // Capturing 'self' weak, so we have to pass in a metatype explicitly
+    // so that IRGen can recover metadata.
+
+    // CHECK:      [[WEAK_SELF:%.*]] = alloc_box $@sil_weak Optional<Z>
+    // CHECK:      [[FN:%.*]] = function_ref @_TFFC12dynamic_self1Z23testDynamicSelfCapturesFT1xSi_DS0_U1_FT_T_ : $@convention(thin) (@owned @box @sil_weak Optional<Z>, @thick Z.Type) -> ()
+    // CHECK:      strong_retain [[WEAK_SELF]] : $@box @sil_weak Optional<Z>
+    // CHECK-NEXT: [[DYNAMIC_SELF:%.*]] = metatype $@thick Self.Type
+    // CHECK-NEXT: [[STATIC_SELF:%.*]] = upcast [[DYNAMIC_SELF]] : $@thick Self.Type to $@thick Z.Type
+    // CHECK:      partial_apply [[FN]]([[WEAK_SELF]], [[STATIC_SELF]]) : $@convention(thin) (@owned @box @sil_weak Optional<Z>, @thick Z.Type) -> ()
+    let fn3 = {
+      [weak self] in
+      _ = self
+    }
+    fn3()
+
+    return self
+  }
+
+}
+
 // CHECK-LABEL: sil_witness_table hidden X: P module dynamic_self {
 // CHECK: method #P.f!1: @_TTWC12dynamic_self1XS_1PS_FS1_1f
 

test/SILOptimizer/inline_self.swift

@@ -2,6 +2,8 @@
 //
 // This is a .swift test because the SIL parser does not support Self.
 
+// Do not inline C.factory into main. Doing so would lose the ability
+// to materialize local Self metadata.
 class C {
   required init() {}
 }
@@ -22,14 +24,32 @@ extension C {
   }
 }
 
-// Do not inline C.factory into main. Doing so would lose the ability
-// to materialize local Self metadata.
-//
+// Call the function so it can be inlined.
+var x = C()
+var x2 = C.factory(1)
+
+@inline(never)
+func callIt(fn: () -> ()) {
+  fn()
+}
+
+// Do not inline C.capturesSelf() into main either.
+class Z {
+  func capturesSelf() -> Self {
+    let fn = { [weak self] in _ = self }
+    callIt(fn: fn)
+    return self
+  }
+}
+
+_ = Z().capturesSelf()
+
 // CHECK-LABEL: sil @main : $@convention(c)
 // CHECK: function_ref static inline_self.C.factory (Swift.Int) -> Self
 // CHECK: [[F:%[0-9]+]] = function_ref @_TZFC11inline_self1C7factory{{.*}} : $@convention(method) (Int, @thick C.Type) -> @owned C
 // CHECK: apply [[F]](%{{.+}}, %{{.+}}) : $@convention(method) (Int, @thick C.Type) -> @owned C
 
-// Call the function so it can be inlined.
-var x = C()
-var x2 = C.factory(1)
+// CHECK: [[Z:%.*]] = alloc_ref $Z
+// CHECK: function_ref inline_self.Z.capturesSelf () -> Self
+// CHECK: [[F:%[0-9]+]] = function_ref @_TFC11inline_self1Z12capturesSelffT_DS0_ : $@convention(method) (@guaranteed Z) -> @owned Z
+// CHECK: apply [[F]]([[Z]]) : $@convention(method) (@guaranteed Z) -> @owned Z