Merge pull request #4886 from aschwaighofer/fix_irgen_partial_apply_of_generic_witness_swift3

IRGen: Shuffle self value for witness method partial applications

Author: tkremenek

lib/IRGen/GenFunc.cpp

@@ -841,6 +841,13 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
       (origType->hasSelfParam() &&
        isSelfContextParameter(origType->getSelfParameter()));
 
+  // Witness method calls expect self, followed by the self type followed by,
+  // the witness table at the end of the parameter list. But polymorphic
+  // arguments come before this.
+  bool isWitnessMethodCallee = origType->getRepresentation() ==
+      SILFunctionTypeRepresentation::WitnessMethod;
+  Explosion witnessMethodSelfValue;
+
   // If there's a data pointer required, but it's a swift-retainable
   // value being passed as the context, just forward it down.
   if (!layout) {
@@ -993,9 +1000,11 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
         emitApplyArgument(subIGF, origParamInfo,
                           substType->getParameters()[origParamI],
                           param, origParam);
-
-        needsAllocas |=
-          addNativeArgument(subIGF, origParam, origParamInfo, args);
+        bool isWitnessMethodCalleeSelf = (isWitnessMethodCallee &&
+            origParamI + 1 == origType->getParameters().size());
+        needsAllocas |= addNativeArgument(
+            subIGF, origParam, origParamInfo,
+            isWitnessMethodCalleeSelf ? witnessMethodSelfValue : args);
         ++origParamI;
       } else {
         args.add(param.claimAll());
@@ -1043,8 +1052,7 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
   // witness table. Metadata for Self is derived inside the partial
   // application thunk and doesn't need to be stored in the outer
   // context.
-  if (origType->getRepresentation() ==
-      SILFunctionTypeRepresentation::WitnessMethod) {
+  if (isWitnessMethodCallee) {
     assert(fnContext->getType() == IGM.Int8PtrTy);
     llvm::Value *wtable = subIGF.Builder.CreateBitCast(
         fnContext, IGM.WitnessTablePtrTy);
@@ -1061,6 +1069,11 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
     args.add(llvm::UndefValue::get(IGM.RefCountedPtrTy));
   }
 
+  // Add the witness methods self argument before the error parameter after the
+  // polymorphic arguments.
+  if (isWitnessMethodCallee)
+    witnessMethodSelfValue.transferInto(args, witnessMethodSelfValue.size());
+
   // Pass down the error result.
   if (origType->hasErrorResult()) {
     llvm::Value *errorResultPtr = origParams.claimNext();
@@ -1070,8 +1083,7 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
 
   assert(origParams.empty());
 
-  if (origType->getRepresentation() ==
-      SILFunctionTypeRepresentation::WitnessMethod) {
+  if (isWitnessMethodCallee) {
     assert(witnessMetadata.SelfMetadata->getType() == IGM.TypeMetadataPtrTy);
     args.add(witnessMetadata.SelfMetadata);
     assert(witnessMetadata.SelfWitnessTable->getType() == IGM.WitnessTablePtrTy);

test/IRGen/partial_apply_forwarder.sil

@@ -6,6 +6,22 @@ public class C : P {}
 public class D<T : P> {}
 class E {}
 
+public protocol Observable {
+  associatedtype Result
+  func subscribe<T: Observer>(o: T) -> ()
+}
+
+public protocol Observer {
+  associatedtype Result
+}
+
+sil hidden @witness_method : $@convention(thin) <S where S : Observable><O where O : Observer, O.Result == S.Result> (@in S) -> @owned @callee_owned (@in O) -> () {
+bb0(%0 : $*S):
+  %1 = witness_method $S, #Observable.subscribe!1 : $@convention(witness_method) <τ_0_0 where τ_0_0 : Observable><τ_1_0 where τ_1_0 : Observer, τ_1_0.Result == τ_0_0.Result> (@in τ_1_0, @in_guaranteed τ_0_0) -> ()
+  %2 = partial_apply %1<S, S.Result, O>(%0) : $@convention(witness_method) <τ_0_0 where τ_0_0 : Observable><τ_1_0 where τ_1_0 : Observer, τ_1_0.Result == τ_0_0.Result> (@in τ_1_0, @in_guaranteed τ_0_0) -> ()
+  return %2 : $@callee_owned (@in O) -> ()
+}
+
 // CHECK-LABEL: define internal %GC23partial_apply_forwarder1DCS_1C_* @_TPA_unspecialized_uncurried(%swift.refcounted*
 // CHECK: [[TYPE:%.*]] = call %swift.type* @_TMaC23partial_apply_forwarder1C()
 // CHECK: [[CAST:%.*]] = bitcast %swift.refcounted* %0 to %C23partial_apply_forwarder1E*