We need to be very careful in generic specialization of recurisive fu…

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -49,43 +49,59 @@ class GenericSpecializer : public SILFunctionTransform {
 
 bool GenericSpecializer::specializeAppliesInFunction(SILFunction &F) {
   DeadInstructionSet DeadApplies;
+  llvm::SmallSetVector<SILInstruction *, 8> Applies;
 
+  bool Changed = false;
   for (auto &BB : F) {
-    for (auto It = BB.begin(), End = BB.end(); It != End;) {
-      auto &I = *It++;
+    // Collect the applies for this block in reverse order so that we
+    // can pop them off the end of our vector and process them in
+    // forward order.
+    for (auto It = BB.rbegin(), End = BB.rend(); It != End; ++It) {
+      auto *I = &*It;
 
       // Skip non-apply instructions, apply instructions with no
       // substitutions, apply instructions where we do not statically
       // know the called function, and apply instructions where we do
       // not have the body of the called function.
-
-      ApplySite Apply = ApplySite::isa(&I);
+      ApplySite Apply = ApplySite::isa(I);
       if (!Apply || !Apply.hasSubstitutions())
         continue;
 
       auto *Callee = Apply.getReferencedFunction();
       if (!Callee || !Callee->isDefinition())
         continue;
 
-      // Do not attempt to specialize known dead instructions. Doing
-      // so would be a waste of time (since they are unused), and can
-      // also lead to verification errors on the newly created
-      // apply. This can happen in the case of a partial application
-      // of a reabstraction thunk where we have done an RAUW on the
-      // reabstracted function (which is an argument of the partial
-      // apply). In this case we add the partial apply of the
-      // reabstraction thunk to the set of dead applies, but its
-      // arguments types do not match the expected types of the
-      // argument that has been RAUWed into it.
-      if (DeadApplies.count(Apply.getInstruction()))
-        continue;
+      Applies.insert(Apply.getInstruction());
+    }
+
+    // Attempt to specialize each apply we collected, deleting any
+    // that we do specialize (along with other instructions we clone
+    // in the process of doing so). We pop from the end of the list to
+    // avoid tricky iterator invalidation issues.
+    while (!Applies.empty()) {
+      auto *I = Applies.pop_back_val();
+      auto Apply = ApplySite::isa(I);
+      assert(Apply && "Expected an apply!");
 
       // We have a call that can potentially be specialized, so
       // attempt to do so.
-
       llvm::SmallVector<SILFunction *, 2> NewFunctions;
       trySpecializeApplyOfGeneric(Apply, DeadApplies, NewFunctions);
 
+      // Remove all the now-dead applies. We must do this immediately
+      // rather than defer it in order to avoid problems with cloning
+      // dead instructions when doing recursive specialization.
+      while (!DeadApplies.empty()) {
+        auto *AI = DeadApplies.pop_back_val();
+
+        // Remove any applies we are deleting so that we don't attempt
+        // to specialize them.
+        Applies.remove(AI);
+
+        recursivelyDeleteTriviallyDeadInstructions(AI, true);
+        Changed = true;
+      }
+
       // If calling the specialization utility resulted in new functions
       // (as opposed to returning a previous specialization), we need to notify
       // the pass manager so that the new functions get optimized.
@@ -95,14 +111,6 @@ bool GenericSpecializer::specializeAppliesInFunction(SILFunction &F) {
     }
   }
 
-  // Remove all the now-dead applies.
-  bool Changed = false;
-  while (!DeadApplies.empty()) {
-    auto *AI = DeadApplies.pop_back_val();
-    recursivelyDeleteTriviallyDeadInstructions(AI, true);
-    Changed = true;
-  }
-
   return Changed;
 }
 

test/SILOptimizer/specialize_recursive_generics.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt -enable-sil-verify-all %s -generic-specializer | FileCheck %s
+// RUN: %target-sil-opt -enable-sil-verify-all %s -generic-specializer -cse | FileCheck %s
 
 // Check that SIL cloner can correctly handle specialization of recursive
 // functions with generic arguments.
@@ -89,3 +89,61 @@ bb0:
   %13 = tuple ()                                  // user: %14
   return %13 : $()                                // id: %14
 }
+
+
+public class C : P {}
+
+public protocol P {}
+
+sil hidden [noinline] @helper : $@convention(thin) <T> (@in T, @in P) -> @owned Optional<C> {
+bb0(%0 : $*T, %1 : $*P):
+  %4 = alloc_stack $P
+  copy_addr %1 to [initialization] %4 : $*P
+  %6 = alloc_stack $C
+  checked_cast_addr_br take_always P in %4 : $*P to C in %6 : $*C, bb1, bb2
+bb1:
+  %8 = load %6 : $*C
+  %9 = enum $Optional<C>, #Optional.some!enumelt.1, %8 : $C
+  dealloc_stack %6 : $*C
+  br bb3(%9 : $Optional<C>)
+bb2:
+  %12 = enum $Optional<C>, #Optional.none!enumelt
+  dealloc_stack %6 : $*C
+  br bb3(%12 : $Optional<C>)
+
+bb3(%15 : $Optional<C>):
+  dealloc_stack %4 : $*P
+  destroy_addr %1 : $*P
+  destroy_addr %0 : $*T
+  return %15 : $Optional<C>
+}
+
+// CHECK-LABEL: sil shared @_TTSg5C4main1CS0_S_1PS___lookup
+sil @lookup : $@convention(method) <Self where Self : P> (@owned C, @in_guaranteed Self) -> @owned Optional<C> {
+bb0(%0 : $C, %1 : $*Self):
+  // CHECK: [[HELPER:%.*]] = function_ref @_TTSg5Vs5Int32__helper
+  %4 = function_ref @helper : $@convention(thin) <τ_0_0> (@in τ_0_0, @in P) -> @owned Optional<C>
+  %5 = integer_literal $Builtin.Int32, 1
+  %6 = struct $Int32 (%5 : $Builtin.Int32)
+  %7 = alloc_stack $Int32
+  store %6 to %7 : $*Int32
+  %9 = alloc_stack $P
+  %10 = init_existential_addr %9 : $*P, $Self
+  copy_addr %1 to [initialization] %10 : $*Self
+  // CHECK: apply [[HELPER]]
+  // CHECK-NOT: apply [[HELPER]]
+  %12 = apply %4<Int32>(%7, %9) : $@convention(thin) <τ_0_0> (@in τ_0_0, @in P) -> @owned Optional<C>
+  release_value %12 : $Optional<C>
+  dealloc_stack %9 : $*P
+  dealloc_stack %7 : $*Int32
+  // CHECK: [[LOOKUP:%.*]] = function_ref @_TTSg5C4main1CS0_S_1PS___lookup
+  %16 = function_ref @lookup : $@convention(method) <τ_0_0 where τ_0_0 : P> (@owned C, @in_guaranteed τ_0_0) -> @owned Optional<C>
+  %17 = alloc_stack $C
+  store %0 to %17 : $*C
+  strong_retain %0 : $C
+  // CHECK: apply [[LOOKUP]]
+  %20 = apply %16<C>(%0, %17) : $@convention(method) <τ_0_0 where τ_0_0 : P> (@owned C, @in_guaranteed τ_0_0) -> @owned Optional<C>
+  dealloc_stack %17 : $*C
+  strong_release %0 : $C
+  return %20 : $Optional<C>
+}