IRGen: Fix specialized conformances with abstract conditional requirements in the partial apply forwarder

We need to add abstract requirements of the conditional requirements.

rdar://59456064

Author: aschwaighofer

lib/IRGen/GenFunc.cpp

@@ -1292,8 +1292,9 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
   SmallVector<ParameterConvention, 4> argConventions;
 
   // Reserve space for polymorphic bindings.
-  auto bindings = NecessaryBindings::forFunctionInvocations(IGF.IGM,
-                                                            origType, subs);
+  auto bindings =
+      NecessaryBindings::forPartialApplyForwarder(IGF.IGM, origType, subs);
+
   if (!bindings.empty()) {
     hasSingleSwiftRefcountedContext = No;
     auto bindingsSize = bindings.getBufferSize(IGF.IGM);

lib/IRGen/GenProto.cpp

@@ -2732,9 +2732,43 @@ void NecessaryBindings::addTypeMetadata(CanType type) {
   Requirements.insert({type, nullptr});
 }
 
+/// Add all the abstract conditional conformances in the specialized
+/// conformance to the \p requirements.
+static void addAbstractConditionalRequirements(
+    SpecializedProtocolConformance *specializedConformance,
+    llvm::SetVector<GenericRequirement> &requirements) {
+  auto module = specializedConformance->getDeclContext()->getParentModule();
+  auto subMap = specializedConformance->getSubstitutions(module);
+  auto condRequirements =
+      specializedConformance->getConditionalRequirementsIfAvailable();
+  if (!condRequirements)
+    return;
+
+  for (auto req : *condRequirements) {
+    if (req.getKind() != RequirementKind::Conformance)
+      continue;
+    auto *proto =
+        req.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto ty = req.getFirstType()->getCanonicalType();
+    auto conformance = subMap.lookupConformance(ty, proto);
+    if (!conformance.isAbstract())
+      continue;
+    requirements.insert({ty, conformance.getAbstract()});
+  }
+}
+
 void NecessaryBindings::addProtocolConformance(CanType type,
                                                ProtocolConformanceRef conf) {
-  if (!conf.isAbstract()) return;
+  if (!conf.isAbstract()) {
+    auto specializedConf =
+        dyn_cast<SpecializedProtocolConformance>(conf.getConcrete());
+    // The partial apply forwarder does not have the context to reconstruct
+    // abstract conditional conformance requirements.
+    if (forPartialApply && specializedConf) {
+      addAbstractConditionalRequirements(specializedConf, Requirements);
+    }
+    return;
+  }
   assert(isa<ArchetypeType>(type));
 
   // TODO: pass something about the root conformance necessary to
@@ -2925,7 +2959,24 @@ NecessaryBindings
 NecessaryBindings::forFunctionInvocations(IRGenModule &IGM,
                                           CanSILFunctionType origType,
                                           SubstitutionMap subs) {
+  return computeBindings(IGM, origType, subs,
+                         false /*forPartialApplyForwarder*/);
+}
+
+NecessaryBindings
+NecessaryBindings::forPartialApplyForwarder(IRGenModule &IGM,
+                                          CanSILFunctionType origType,
+                                          SubstitutionMap subs) {
+  return computeBindings(IGM, origType, subs,
+                         true /*forPartialApplyForwarder*/);
+}
+
+NecessaryBindings NecessaryBindings::computeBindings(
+    IRGenModule &IGM, CanSILFunctionType origType, SubstitutionMap subs,
+    bool forPartialApplyForwarder) {
+
   NecessaryBindings bindings;
+  bindings.forPartialApply = forPartialApplyForwarder;
 
   // Bail out early if we don't have polymorphic parameters.
   if (!hasPolymorphicParameters(origType))

lib/IRGen/NecessaryBindings.h

@@ -42,6 +42,13 @@ namespace irgen {
 class NecessaryBindings {
   llvm::SetVector<GenericRequirement> Requirements;
 
+
+  /// Are the bindings to be computed for a partial apply forwarder.
+  /// In the case this is true we need to store/restore the conformance of a
+  /// specialized type with conditional conformance because the conditional
+  /// requirements are not available in the partial apply forwarder.
+  bool forPartialApply = false;
+
 public:
   NecessaryBindings() = default;
 
@@ -50,7 +57,10 @@ class NecessaryBindings {
   static NecessaryBindings forFunctionInvocations(IRGenModule &IGM,
                                                   CanSILFunctionType origType,
                                                   SubstitutionMap subs);
-  
+  static NecessaryBindings forPartialApplyForwarder(IRGenModule &IGM,
+                                                    CanSILFunctionType origType,
+                                                    SubstitutionMap subs);
+
   /// Add whatever information is necessary to reconstruct type metadata
   /// for the given type.
   void addTypeMetadata(CanType type);
@@ -84,6 +94,11 @@ class NecessaryBindings {
   const llvm::SetVector<GenericRequirement> &getRequirements() const {
     return Requirements;
   }
+private:
+  static NecessaryBindings computeBindings(IRGenModule &IGM,
+                                           CanSILFunctionType origType,
+                                           SubstitutionMap subs,
+                                           bool forPartialApplyForwarder);
 };
 
 } // end namespace irgen

test/IRGen/partial_apply_forwarder.sil

@@ -205,7 +205,7 @@ sil hidden_external @takingQAndS : $@convention(thin) <τ_0_0 where  τ_0_0 : Q>
 // CHECK:   [[WBREF:%.*]] = bitcast %T23partial_apply_forwarder7WeakBoxC* [[WB]] to %swift.refcounted*
 // CHECK:   call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[WBREF]])
 // CHECK:   [[TYADDR:%.*]] = getelementptr inbounds %T23partial_apply_forwarder7WeakBoxC, %T23partial_apply_forwarder7WeakBoxC* [[WB]], i32 0, i32 0, i32 0
-// CHECK:   [[TY:%.*]] = load %swift.type*, %swift.type** [[TYADDR]] 
+// CHECK:   [[TY:%.*]] = load %swift.type*, %swift.type** [[TYADDR]]
 // CHECK:   %.asUnsubstituted = bitcast %T23partial_apply_forwarder7WeakBoxC* [[WB]] to %T23partial_apply_forwarder7WeakBoxC.1*
 // CHECK:   call void @swift_release(%swift.refcounted* %0)
 // CHECK:   [[GENPARAMSADDR:%.*]] = bitcast %swift.type* [[TY]] to %swift.type**
@@ -248,6 +248,42 @@ bb0:
 // CHECK: entry:
 // CHECK: ret { i8*, %swift.refcounted* } { i8* bitcast ({ %{{.*}}Empty{{.*}}*, i8*, %swift.refcounted* } (%{{.*}}Empty{{.*}}*, %swift.refcounted*)* @"{{.*}}returns_closure{{.*}}" to i8*), %swift.refcounted* null }
 
+protocol MyEquatable {
+  static func isEqual (lhs: Self, rhs: Self) -> Builtin.Int1
+}
+
+public struct Inner<Element> {
+  public init()
+  public init(_ e: Element)
+}
+
+extension Inner : MyEquatable where Element : MyEquatable {
+  public static func isEqual (lhs: Inner<Element>, rhs: Inner<Element>) -> Builtin.Int1
+}
+
+public struct Outer<Value> {
+  init()
+}
+
+extension Outer : MyEquatable where Value : MyEquatable {
+  public static func isEqual (lhs: Outer<Value>, rhs: Outer<Value>) -> Builtin.Int1
+}
+
+sil @$closure : $@convention(method) <Value where Value : MyEquatable> (Outer<Value>, Outer<Value>, @thin Outer<Value>.Type) -> Builtin.Int1
+
+sil @$dont_crash_test_caputre_specialized_conditional_conformance : $@convention(thin) <Element where Element : MyEquatable> (Outer<Inner<Element>>) -> () {
+bb0(%0 : $Outer<Inner<Element>>):
+  %2 = alloc_stack $Outer<Inner<Element>>
+  store %0 to %2 : $*Outer<Inner<Element>>
+  %4 = metatype $@thin Outer<Inner<Element>>.Type
+  %5 = function_ref @$closure : $@convention(method) <τ_0_0 where τ_0_0 : MyEquatable> (Outer<τ_0_0>, Outer<τ_0_0>, @thin Outer<τ_0_0>.Type) -> Builtin.Int1
+  %6 = partial_apply [callee_guaranteed] %5<Inner<Element>>(%4) : $@convention(method) <τ_0_0 where τ_0_0 : MyEquatable> (Outer<τ_0_0>, Outer<τ_0_0>, @thin Outer<τ_0_0>.Type) -> Builtin.Int1
+  strong_release %6 : $@callee_guaranteed (Outer<Inner<Element>>, Outer<Inner<Element>>) -> Builtin.Int1
+  dealloc_stack %2 : $*Outer<Inner<Element>>
+  %15 = tuple ()
+  return %15 : $()
+}
+
 sil_vtable WeakBox {}
 sil_vtable C {}
 sil_vtable D {}