IRGen: Restore handling of "simple" partial_apply instructions.

Methods and closures use the same convention for the self/context argument, so when a
partial_apply applies a single argument to a method implementation, the closure can be formed
by simply tupling the implementation and self argument together. This doesn't yet cover a lot
of situations the compiler generates, but it prepares IRGen for an IRGen SIL pass that can
take the heavy lifting of partial_apply lowering off of IRGen by reducing more complex closures
into these "simple" cases, eventually reducing the number of partial_apply forwarding thunks and
extra closure allocations generated.

Author: jckarter

lib/IRGen/IRGenSIL.cpp

@@ -2589,9 +2589,107 @@ getPartialApplicationFunction(IRGenSILFunction &IGF, SILValue v,
   llvm_unreachable("bad kind");
 }
 
+// A "simple" partial_apply is one where the argument can be directly
+// adopted as the context of the result closure.
+static bool isSimplePartialApply(IRGenFunction &IGF, PartialApplyInst *i) {
+  // The callee type must use the `method` convention.
+  auto calleeTy = i->getCallee()->getType().castTo<SILFunctionType>();
+  auto resultTy = i->getFunctionType();
+  
+  if (calleeTy->getRepresentation() != SILFunctionTypeRepresentation::Method)
+    return false;
+  
+  // There should be one applied argument.
+  // (This is a bit stricter than necessary, because empty arguments could be
+  // ignored, and for noescape closures, any amount of data less than a pointer
+  // in size can be blobbed into a single context word, but those will be
+  // handled by a simplification pass in SIL.)
+  if (i->getNumArguments() != 1)
+    return false;
+  
+  auto appliedParam = calleeTy->getParameters().back();
+  if (resultTy->isNoEscape()) {
+    // A trivial closure accepts an unowned or guaranteed argument, possibly
+    // direct or indirect.
+    switch (appliedParam.getConvention()) {
+    case ParameterConvention::Indirect_Inout:
+    case ParameterConvention::Indirect_In_Constant:
+    case ParameterConvention::Indirect_In_Guaranteed:
+    case ParameterConvention::Indirect_InoutAliasable:
+      // Indirect arguments are trivially word sized.
+      return true;
+        
+    case ParameterConvention::Direct_Guaranteed:
+    case ParameterConvention::Direct_Unowned: {
+      // Is the direct argument a single word-sized value?
+      auto argSchema = IGF.IGM.getTypeInfo(i->getArgument(0)->getType())
+                          .getSchema();
+      if (argSchema.size() != 1)
+        return false;
+        
+      if (argSchema[0].getScalarType()->getPrimitiveSizeInBits()
+            != IGF.IGM.getPointerSize().getValueInBits())
+        return false;
+      
+      return true;
+    }
+    default:
+      return false;
+    }
+  } else {
+    // An escaping closure argument's convention should match the callee
+    // convention of the result.
+    if (resultTy->getCalleeConvention() != appliedParam.getConvention()) {
+      return false;
+    }
+    assert(!isIndirectFormalParameter(resultTy->getCalleeConvention()));
+    
+    auto &argInfo = IGF.IGM.getTypeInfo(i->getArgument(0)->getType());
+    
+    if (!argInfo.isSingleSwiftRetainablePointer(ResilienceExpansion::Maximal))
+      return false;
+    
+    return true;
+  }
+}
+
 void IRGenSILFunction::visitPartialApplyInst(swift::PartialApplyInst *i) {
   SILValue v(i);
 
+  if (isSimplePartialApply(*this, i)) {
+    Explosion function;
+    
+    auto schema = IGM.getTypeInfo(v->getType()).getSchema();
+    assert(schema.size() == 2);
+    auto calleeTy = schema[0].getScalarType();
+    auto contextTy = schema[1].getScalarType();
+    
+    auto callee = getLoweredExplosion(i->getCallee());
+    auto calleeValue = callee.claimNext();
+    assert(callee.empty());
+    calleeValue = Builder.CreateBitOrPointerCast(calleeValue, calleeTy);
+    function.add(calleeValue);
+
+    Explosion context;
+    for (auto arg : i->getArguments()) {
+      auto &value = getLoweredValue(arg);
+      
+      if (value.isAddress()) {
+        context.add(value.getAnyAddress().getAddress());
+      } else {
+        getLoweredExplosion(arg, context);
+      }
+    }
+    auto contextValue = context.claimNext();
+    assert(context.empty());
+    contextValue = Builder.CreateBitOrPointerCast(contextValue, contextTy);
+    function.add(contextValue);
+    
+    setLoweredExplosion(v, function);
+    return;
+  }
+  
+
   // NB: We collect the arguments under the substituted type.
   auto args = i->getArguments();
   auto calleeTy = i->getSubstCalleeType();

test/IRGen/simple_partial_apply.sil

@@ -0,0 +1,57 @@
+// RUN: %target-swift-frontend -emit-ir %s | %FileCheck %s
+
+sil_stage canonical
+
+import Swift
+
+class C {}
+
+sil_vtable C {}
+
+struct SingleRefcounted {
+    var c: C
+}
+
+// CHECK-LABEL: define {{.*}} @escape_partial_apply_swift_class
+// CHECK:      [[FPTR:%.*]] = insertvalue { i8*, %swift.refcounted* } undef, i8* %0, 0
+// CHECK-NEXT: [[FCTX:%.*]] = insertvalue { i8*, %swift.refcounted* } [[FPTR]], %swift.refcounted* {{.*}}, 1
+// CHECK-NEXT: ret { i8*, %swift.refcounted* } [[FCTX]]
+sil @escape_partial_apply_swift_class : $@convention(thin) (@convention(method) (Int, @guaranteed C) -> Int, @guaranteed C) -> @callee_guaranteed (Int) -> Int {
+entry(%body : $@convention(method) (Int, @guaranteed C) -> Int, %context : $C):
+  %closure = partial_apply [callee_guaranteed] %body(%context) : $@convention(method) (Int, @guaranteed C) -> Int
+  return %closure : $@callee_guaranteed (Int) -> Int
+}
+
+// CHECK-LABEL: define {{.*}} @escape_partial_apply_swift_single_refcount_struct
+// CHECK:      [[FPTR:%.*]] = insertvalue { i8*, %swift.refcounted* } undef, i8* %0, 0
+// CHECK-NEXT: [[FCTX:%.*]] = insertvalue { i8*, %swift.refcounted* } [[FPTR]], %swift.refcounted* {{.*}}, 1
+// CHECK-NEXT: ret { i8*, %swift.refcounted* } [[FCTX]]
+sil @escape_partial_apply_swift_single_refcount_struct : $@convention(thin) (@convention(method) (Int, @guaranteed SingleRefcounted) -> Int, @guaranteed SingleRefcounted) -> @callee_guaranteed (Int) -> Int {
+entry(%body : $@convention(method) (Int, @guaranteed SingleRefcounted) -> Int, %context : $SingleRefcounted):
+  %closure = partial_apply [callee_guaranteed] %body(%context) : $@convention(method) (Int, @guaranteed SingleRefcounted) -> Int
+  return %closure : $@callee_guaranteed (Int) -> Int
+}
+
+// CHECK-LABEL: define {{.*}} @noescape_partial_apply_swift_indirect
+// CHECK:         [[CTX:%.*]] = bitcast {{.*}}** %1 to %swift.opaque*
+// CHECK-NEXT:    [[CONT:%.*]] = bitcast i8* %2
+// CHECK-NEXT:    call {{.*}}void [[CONT]](i8* %0, %swift.opaque* [[CTX]], %swift.refcounted* {{.*}}%3)
+sil @noescape_partial_apply_swift_indirect : $@convention(thin) (@convention(method) (Int, @in_guaranteed C) -> Int, @in_guaranteed C, @guaranteed @callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()) -> () {
+entry(%body : $@convention(method) (Int, @in_guaranteed C) -> Int, %context : $*C, %cont : $@callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()):
+  %closure = partial_apply [callee_guaranteed] [on_stack] %body(%context) : $@convention(method) (Int, @in_guaranteed C) -> Int
+  %x = apply %cont(%closure) : $@callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()
+  dealloc_stack %closure : $@noescape @callee_guaranteed (Int) -> Int
+  return undef : $()
+}
+
+// CHECK-LABEL: define {{.*}} @noescape_partial_apply_swift_direct_word
+// CHECK:         [[CTX:%.*]] = inttoptr i{{.*}} %1 to %swift.opaque*
+// CHECK-NEXT:    [[CONT:%.*]] = bitcast i8* %2
+// CHECK-NEXT:    call {{.*}}void [[CONT]](i8* %0, %swift.opaque* [[CTX]], %swift.refcounted* {{.*}}%3)
+sil @noescape_partial_apply_swift_direct_word : $@convention(thin) (@convention(method) (Int, Int) -> Int, Int, @guaranteed @callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()) -> () {
+entry(%body : $@convention(method) (Int, Int) -> Int, %context : $Int, %cont : $@callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()):
+  %closure = partial_apply [callee_guaranteed] [on_stack] %body(%context) : $@convention(method) (Int, Int) -> Int
+  %x = apply %cont(%closure) : $@callee_guaranteed (@noescape @callee_guaranteed (Int) -> Int) -> ()
+  dealloc_stack %closure : $@noescape @callee_guaranteed (Int) -> Int
+  return undef : $()
+}