UnsafeGuaranteedPeephole: Also use RCIdentity for matching retains

Author: aschwaighofer

lib/SILOptimizer/Transforms/UnsafeGuaranteedPeephole.cpp

@@ -129,6 +129,7 @@ static SILBasicBlock::iterator findReleaseToMatchUnsafeGuaranteedValue(
 /// sequences.
 static bool removeGuaranteedRetainReleasePairs(SILFunction &F,
                                                RCIdentityFunctionInfo &RCIA) {
+  DEBUG(llvm::dbgs() << "Running on function " << F.getName() << "\n");
   bool Changed = false;
   for (auto &BB : F) {
     auto It = BB.begin(), End = BB.end();
@@ -139,7 +140,7 @@ static bool removeGuaranteedRetainReleasePairs(SILFunction &F,
 
       // Memorize the last retain.
       if (isa<StrongRetainInst>(CurInst) || isa<RetainValueInst>(CurInst)) {
-        LastRetain[CurInst->getOperand(0)] = CurInst;
+        LastRetain[RCIA.getRCIdentityRoot(CurInst->getOperand(0))] = CurInst;
         continue;
       }
 
@@ -151,14 +152,15 @@ static bool removeGuaranteedRetainReleasePairs(SILFunction &F,
         continue;
 
       auto Opd = UnsafeGuaranteedI->getOperand(0);
-      if (!LastRetain.count(Opd)) {
+      auto RCIdOpd = RCIA.getRCIdentityRoot(UnsafeGuaranteedI->getOperand(0));
+      if (!LastRetain.count(RCIdOpd)) {
         DEBUG(llvm::dbgs() << "LastRetain failed\n");
         continue;
       }
 
       // This code is very conservative. Check that there is a matching retain
       // before the unsafeGuaranteed builtin with only retains inbetween.
-      auto *LastRetainInst = LastRetain[Opd];
+      auto *LastRetainInst = LastRetain[RCIdOpd];
       auto NextInstIter = std::next(SILBasicBlock::iterator(LastRetainInst));
       while (NextInstIter != BB.end() && &*NextInstIter != CurInst &&
              (isa<RetainValueInst>(*NextInstIter) ||

test/SILOptimizer/unsafe_guaranteed_peephole.sil

@@ -402,3 +402,36 @@ bb0(%0 : $Foo, %1: $Builtin.Int32):
   %17 = tuple ()
   return %17 : $()
 }
+
+// CHECK-LABEL: sil @testUnsafeGuaranteed_cast2_inst
+// CHECK: bb0([[P:%.*]] : $AnyObject
+// CHECK:  [[C:%.*]] = unchecked_ref_cast [[P]] : $AnyObject to $Foo
+// CHECK-NOT: retain
+// CHECK-NOT: unsafeGuaranteed
+// CHECK:  [[U:%.*]] = upcast [[C]] : $Foo to $Base
+// CHECK:   [[M:%.*]] = class_method [[U]] : $Base, #Base.beep
+// CHECK:   apply [[M]]([[U]])
+// CHECK-NOT: release
+// CHECK-NOT: unsafeGuaranteedEnd
+// CHECK:   [[T:%.*]] = tuple ()
+// CHECK:   return [[T]]
+// CHECK: }
+sil @testUnsafeGuaranteed_cast2_inst : $@convention(method) (@guaranteed AnyObject, Builtin.Int32) -> () {
+bb0(%0 : $AnyObject, %1: $Builtin.Int32):
+  %2 = unchecked_ref_cast %0 : $AnyObject to $Foo
+  strong_retain %0 : $AnyObject
+  debug_value %0 : $AnyObject
+  %3 = struct $MyInt(%1 : $Builtin.Int32)
+  %4 = builtin "unsafeGuaranteed"<Foo>(%2 : $Foo) : $(Foo, Builtin.Int8)
+  %5 = tuple_extract %4 : $(Foo, Builtin.Int8), 0
+  %6 = tuple_extract %4 : $(Foo, Builtin.Int8), 1
+  %7 = upcast %5 : $Foo to $Base
+  %19 = class_method %7 : $Base, #Base.beep!1 : (Base) -> () -> () , $@convention(method) (@guaranteed Base) -> ()
+  %20 = apply %19(%7) : $@convention(method) (@guaranteed Base) -> ()
+  strong_release %7 : $Base
+  %21 = struct $MyInt(%1 : $Builtin.Int32)
+  debug_value %5 : $Foo
+  %16 = builtin "unsafeGuaranteedEnd"(%6 : $Builtin.Int8) : $()
+  %17 = tuple ()
+  return %17 : $()
+}