SILOptimizer: Fix speculative devirtualization to correctly use casted value

slavapestov · slavapestov · commit 64971e7e55ec · 2020-01-03T20:17:37.000-05:00
We used to emit this: checked_cast_br [exact] %0 : $Foo to $Bar, bb1, bb2 bb1(%1 : $Bar): %2 = unchecked_ref_cast %0 : $Foo to %1 : $Bar apply ...(..., %2) br ... bb2: ... This is not ownership SIL-safe, because we're re-using the original operand, which will have already been consumed at that point. The more immediate problem here is that this is actually not safe when combined with other optimizations. Suppose that after the speculative devirtualization, our function is inlined into another function where the operand was an upcast. Now we have this code: %0 = upcast ... checked_cast_br [exact] %0 : $Foo to $Bar, bb1, bb2 bb1(%1 : $Bar): %2 = unchecked_ref_cast %0 : $Foo to %1 : $Bar apply ...(..., %2) br ... bb2: ... The SILCombiner will simplify the unchecked_ref_cast of the upcast into an unchecked_ref_cast of the upcast's original operand. At this point, we have an unchecked_ref_cast between two unrelated class types. While this means the block bb1 is unreachable, we might perform further optimizations on it before we run the cast optimizer and delete it altogether. In particular, the devirtualizer follows chains of reference cast instructions, and it will get very confused if it finds this invalid cast between unrelated types. Fixes <rdar://problem/57712094>, <https://bugs.swift.org/browse/SR-11916>.
diff --git a/lib/SILOptimizer/Transforms/SpeculativeDevirtualizer.cpp b/lib/SILOptimizer/Transforms/SpeculativeDevirtualizer.cpp
@@ -81,14 +81,20 @@ static SILBasicBlock *cloneEdge(TermInst *TI, unsigned SuccIndex) {
 }
 
 // A utility function for cloning the apply instruction.
-static FullApplySite CloneApply(FullApplySite AI, SILBuilder &Builder) {
+static FullApplySite CloneApply(FullApplySite AI, SILValue SelfArg,
+                                SILBuilder &Builder) {
   // Clone the Apply.
   Builder.setCurrentDebugScope(AI.getDebugScope());
   Builder.addOpenedArchetypeOperands(AI.getInstruction());
   auto Args = AI.getArguments();
   SmallVector<SILValue, 8> Ret(Args.size());
-  for (unsigned i = 0, e = Args.size(); i != e; ++i)
-    Ret[i] = Args[i];
+  for (unsigned i = 0, e = Args.size(); i != e; ++i) {
+    if (i == e - 1 && SelfArg) {
+      Ret[i] = SelfArg;
+    } else {
+      Ret[i] = Args[i];
+    }
+  }
 
   FullApplySite NAI;
 
@@ -169,8 +175,8 @@ static FullApplySite speculateMonomorphicTarget(FullApplySite AI,
   SILValue DownCastedClassInstance = Iden->getArgument(0);
 
   // Copy the two apply instructions into the two blocks.
-  FullApplySite IdenAI = CloneApply(AI, IdenBuilder);
-  FullApplySite VirtAI = CloneApply(AI, VirtBuilder);
+  FullApplySite IdenAI = CloneApply(AI, DownCastedClassInstance, IdenBuilder);
+  FullApplySite VirtAI = CloneApply(AI, SILValue(), VirtBuilder);
 
   // See if Continue has a release on self as the instruction right after the
   // apply. If it exists, move it into position in the diamond.
diff --git a/test/SILOptimizer/devirt_speculate.swift b/test/SILOptimizer/devirt_speculate.swift
@@ -10,44 +10,96 @@ public class Base {
   public init() {}
   public func foo() {}
 }
+
+@_optimize(none)
+func blackHole<T : AnyObject>(_: T) {}
+
 class Sub1 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub2 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub3 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub4 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub5 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub6 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub7 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 // CHECK: @$s16devirt_speculate28testMaxNumSpeculativeTargetsyyAA4BaseCF
-// CHECK: checked_cast_br [exact] %0 : $Base to Base
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub1
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub2
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub3
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub4
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub5
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub6
+// CHECK: bb0(%0 : $Base):
+// CHECK:   checked_cast_br [exact] %0 : $Base to Base, bb2, bb3
+
+// CHECK: bb2([[CASTED:%.*]]):
+// CHECK:   br bb1
+
+// CHECK: bb3:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub1, bb4, bb5
+
+// CHECK: bb4([[CASTED:%.*]] : $Sub1):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub1>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb5:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub2, bb6, bb7
+
+// CHECK: bb6([[CASTED:%.*]] : $Sub2):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub2>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb7:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub3, bb8, bb9
+
+// CHECK: bb8([[CASTED:%.*]] : $Sub3):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub3>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb9:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub4, bb10, bb11
+
+// CHECK: bb10([[CASTED:%.*]] : $Sub4):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub4>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb11:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub5, bb12, bb13
+
+// CHECK: bb12([[CASTED:%.*]] : $Sub5):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub5>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb13:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub6, bb14, bb15
+
+// CHECK: bb14([[CASTED:%.*]] : $Sub6):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub6>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb15:
 // CHECK-NOT: checked_cast_br
-// CHECK: %[[CM:[0-9]+]] = class_method %0 : $Base, #Base.foo!1 : (Base) -> () -> (), $@convention(method) (@guaranteed Base) -> ()
-// CHECK: apply %[[CM]](%0) : $@convention(method) (@guaranteed Base) -> ()
+// CHECK:   %[[CM:[0-9]+]] = class_method %0 : $Base, #Base.foo!1 : (Base) -> () -> (), $@convention(method) (@guaranteed Base) -> ()
+// CHECK:   apply %[[CM]](%0) : $@convention(method) (@guaranteed Base) -> ()
 
 // YAML:      Pass:            sil-speculative-devirtualizer
 // YAML-NEXT: Name:            sil.PartialSpecDevirt
 // YAML-NEXT: DebugLoc:
 // YAML-NEXT:   File:            {{.*}}/devirt_speculate.swift
-// YAML-NEXT:   Line:            66
+// YAML-NEXT:   Line:            118
 // YAML-NEXT:   Column:          5
 // YAML-NEXT: Function:        'testMaxNumSpeculativeTargets(_:)'
 // YAML-NEXT: Args:
diff --git a/test/SILOptimizer/devirt_speculative.sil b/test/SILOptimizer/devirt_speculative.sil
@@ -145,7 +145,6 @@ bb0(%0: $Base2):
 // CHECK-NEXT: unreachable
 // CHECK: checked_cast_br
 // CHECK: function_ref @Sub_exit
-// CHECK-NEXT: unchecked_ref_cast
 // CHECK-NEXT: apply
 // CHECK-NEXT: unreachable
 sil hidden [noinline] @test_devirt_of_noreturn_function : $@convention(thin) (@owned Base) -> () {
