Fix closure inlining after witness devirtualization.

Certain patterns of directly applied partial_apply's were not being
inlined. This can happen when a closure is defined in the
implementation of a generic witness method. I noticed this issue while
debugging SR-6254: "Fix (or explain) strange ways to make code >3
times faster/slower".

After the witness method is devirtualization and specialized, the
closure application looks like this:

  %pa = partial_apply [callee_guaranteed] %fn() : $@convention(thin) (@in Int) -> @out Int
  %cvt = convert_escape_to_noescape %pa
  apply %cvt(...)

SILCombine already removes the partial_apply and convert_escape_to_noescape generating:

  %thick = thin_to_thick %fn
  apply %thick(...)

However, surprisingly, neither the inliner nor SILCombine can handle this.

I cleaned up the code in SILCombine's apply visitor that handles
thin_to_thick and generalized it to handle this and other patterns.

I also added a restriction to this code so that it doesn't break SIL
ownership in the future, as I discussed with Arnold.

Author: atrick

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -283,6 +283,9 @@ class SILCombiner :
                                                SILValue ConcreteTypeDef,
                                                ProtocolConformanceRef Conformance,
                                                ArchetypeType *OpenedArchetype);
+
+  FullApplySite rewriteApplyCallee(FullApplySite apply, SILValue callee);
+
   SILInstruction *
   propagateConcreteTypeOfInitExistential(FullApplySite AI,
       ProtocolDecl *Protocol,

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -1305,6 +1305,27 @@ bool SILCombiner::optimizeIdentityCastComposition(ApplyInst *FInverse,
   return true;
 }
 
+// Return a new apply with the specified callee. This creates a new apply rather
+// than simply rewriting the callee operand because the apply's SubstCalleeType,
+// derived from the callee and substitution list, may change.
+FullApplySite SILCombiner::rewriteApplyCallee(FullApplySite apply,
+                                              SILValue callee) {
+  SmallVector<SILValue, 4> arguments;
+  for (SILValue arg : apply.getArguments())
+    arguments.push_back(arg);
+
+  Builder.addOpenedArchetypeOperands(apply.getInstruction());
+  if (auto *TAI = dyn_cast<TryApplyInst>(apply)) {
+    return Builder.createTryApply(TAI->getLoc(), callee,
+                                  TAI->getSubstitutions(), arguments,
+                                  TAI->getNormalBB(), TAI->getErrorBB());
+  } else {
+    return Builder.createApply(apply.getLoc(), callee, apply.getSubstitutions(),
+                               arguments,
+                               cast<ApplyInst>(apply)->isNonThrowing());
+  }
+}
+
 SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
   Builder.setCurrentDebugScope(AI->getDebugScope());
   // apply{partial_apply(x,y)}(z) -> apply(z,x,y) is triggered
@@ -1346,20 +1367,11 @@ SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
 
   // (apply (thin_to_thick_function f)) to (apply f)
   if (auto *TTTFI = dyn_cast<ThinToThickFunctionInst>(AI->getCallee())) {
-    // TODO: Handle substitutions and indirect results
-    if (AI->hasSubstitutions() || AI->hasIndirectResults())
-      return nullptr;
-    SmallVector<SILValue, 4> Arguments;
-    for (auto &Op : AI->getArgumentOperands()) {
-      Arguments.push_back(Op.get());
-    }
-    // The type of the substitution is the source type of the thin to thick
-    // instruction.
-    Builder.addOpenedArchetypeOperands(AI);
-    auto *NewAI = Builder.createApply(AI->getLoc(), TTTFI->getOperand(),
-                                      AI->getSubstitutions(), Arguments,
-                                      AI->isNonThrowing());
-    return NewAI;
+    // We currently don't remove any possible retain associated with the thick
+    // function when rewriting the callsite. This should be ok because the
+    // ABI normally expects a guaranteed callee.
+    if (!AI->getOrigCalleeType()->isCalleeConsumed())
+      return rewriteApplyCallee(AI, TTTFI->getOperand()).getInstruction();
   }
 
   // (apply (witness_method)) -> propagate information about
@@ -1477,21 +1489,12 @@ SILInstruction *SILCombiner::visitTryApplyInst(TryApplyInst *AI) {
 
   // (try_apply (thin_to_thick_function f)) to (try_apply f)
   if (auto *TTTFI = dyn_cast<ThinToThickFunctionInst>(AI->getCallee())) {
-    // TODO: Handle substitutions and indirect results
-    if (AI->hasSubstitutions() || AI->hasIndirectResults())
-      return nullptr;
-    SmallVector<SILValue, 4> Arguments;
-    for (auto &Op : AI->getArgumentOperands()) {
-      Arguments.push_back(Op.get());
-    }
-    // The type of the substitution is the source type of the thin to thick
-    // instruction.
-    auto *NewAI = Builder.createTryApply(AI->getLoc(), TTTFI->getOperand(),
-                                         AI->getSubstitutions(), Arguments,
-                                         AI->getNormalBB(), AI->getErrorBB());
-    return NewAI;
+    // We currently don't remove any possible retain associated with the thick
+    // function when rewriting the callsite. This should be ok because the
+    // ABI normally expects a guaranteed callee.
+    if (!AI->getOrigCalleeType()->isCalleeConsumed())
+      return rewriteApplyCallee(AI, TTTFI->getOperand()).getInstruction();
   }
-
   // (apply (witness_method)) -> propagate information about
   // a concrete type from init_existential_addr or init_existential_ref.
   if (auto *WMI = dyn_cast<WitnessMethodInst>(AI->getCallee())) {

test/SILOptimizer/sil_combine.sil

@@ -1522,8 +1522,8 @@ bb0:
 sil @eliminate_thin_to_thick_apply : $@convention(thin) () -> () {
 bb0:
   %0 = function_ref @eliminate_dead_thin_to_thick_function_fun : $@convention(thin) () -> ()
-  %1 = thin_to_thick_function %0 : $@convention(thin) () -> () to $@callee_owned () -> ()
-  %2 = apply %1() : $@callee_owned () -> ()
+  %1 = thin_to_thick_function %0 : $@convention(thin) () -> () to $@callee_guaranteed () -> ()
+  %2 = apply %1() : $@callee_guaranteed () -> ()
   return %2 : $()
 }
 

test/SILOptimizer/sil_combine_apply.sil

@@ -517,3 +517,56 @@ bb0:
   %return = tuple ()
   return %return : $()
 }
+
+sil shared [transparent] [thunk] @genericClosure : $@convention(thin) <T> (@in T) -> @out T {
+bb0(%0 : $*T, %1 : $*T):
+  %12 = tuple ()
+  return %12 : $()
+}
+
+// CHECK-LABEL: sil shared @genericThinToThick : $@convention(thin) () -> ()
+// CHECK: [[F:%.*]] = function_ref @genericClosure : $@convention(thin) <τ_0_0> (@in τ_0_0) -> @out τ_0_0
+// CHECK: apply [[F]]<Builtin.Int64>(%{{.*}}, %{{.*}}) : $@convention(thin) <τ_0_0> (@in τ_0_0) -> @out τ_0_0
+// CHECK-LABEL: } // end sil function 'genericThinToThick'
+sil shared @genericThinToThick : $@convention(thin) () -> () {
+bb0:
+  %fn = function_ref @genericClosure : $@convention(thin) <T> (@in T) -> @out T
+  %thick = thin_to_thick_function %fn : $@convention(thin) <T> (@in T) -> @out T to $@noescape @callee_guaranteed <T> (@in T) -> @out T
+  %in = alloc_stack $Builtin.Int64
+  %out = alloc_stack $Builtin.Int64
+  %c3 = integer_literal $Builtin.Int64, 3
+  store %c3 to %in : $*Builtin.Int64
+  %call = apply %thick<Builtin.Int64>(%out, %in) : $@noescape @callee_guaranteed <T> (@in T) -> @out T
+  dealloc_stack %out : $*Builtin.Int64
+  dealloc_stack %in : $*Builtin.Int64
+  %999 = tuple ()
+  return %999 : $()
+}
+
+sil shared [transparent] [thunk] @indirectClosure : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64 {
+bb0(%0 : $*Builtin.Int64, %1 : $*Builtin.Int64):
+  %val = load %1 : $*Builtin.Int64
+  store %val to %0 : $*Builtin.Int64
+  %999 = tuple ()
+  return %999 : $()
+}
+
+// CHECK-LABEL: sil shared @appliedEscapeToNoEscape : $@convention(thin) () -> () {
+// CHECK: [[F:%.*]] = function_ref @indirectClosure : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64 // user: %5
+// CHECK: apply [[F]](%{{.*}}, %{{.*}}) : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+sil shared @appliedEscapeToNoEscape : $@convention(thin) () -> () {
+bb0:
+  %fn = function_ref @indirectClosure : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+  %pa = partial_apply [callee_guaranteed] %fn() : $@convention(thin) (@in Builtin.Int64) -> @out Builtin.Int64
+  %cvt = convert_escape_to_noescape %pa : $@callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64 to $@noescape @callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64
+  %out = alloc_stack $Builtin.Int64
+  %in = alloc_stack $Builtin.Int64
+  %c3 = integer_literal $Builtin.Int64, 3
+  store %c3 to %in : $*Builtin.Int64
+  %call = apply %cvt(%out, %in) : $@noescape @callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64
+  dealloc_stack %in : $*Builtin.Int64
+  dealloc_stack %out : $*Builtin.Int64
+  strong_release %pa : $@callee_guaranteed (@in Builtin.Int64) -> @out Builtin.Int64
+  %999 = tuple ()
+  return %999 : $()
+}