Merge pull request #38864 from meg-gupta/fixinvalidrauwbug

Fix ownership rauw to not leave behind stale ownership fixup context

Author: atrick

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -144,6 +144,11 @@ class OwnershipRAUWHelper {
 private:
   SILBasicBlock::iterator replaceAddressUses(SingleValueInstruction *oldValue,
                                              SILValue newValue);
+
+  void invalidate() {
+    ctx->clear();
+    ctx = nullptr;
+  }
 };
 
 /// A utility composed ontop of OwnershipFixupContext that knows how to replace
@@ -175,6 +180,12 @@ class OwnershipReplaceSingleUseHelper {
 
   /// Perform the actual RAUW.
   SILBasicBlock::iterator perform();
+
+private:
+  void invalidate() {
+    ctx->clear();
+    ctx = nullptr;
+  }
 };
 
 /// An abstraction over LoadInst/LoadBorrowInst so one can handle both types of

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -859,7 +859,7 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
   // Otherwise, lets check if we can perform this RAUW operation. If we can't,
   // set ctx to nullptr to invalidate the helper and return.
   if (!canFixUpOwnershipForRAUW(oldValue, newValue, inputCtx)) {
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
@@ -908,16 +908,14 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
     return;
 
   if (!borrowedAddress.interiorPointerOp) {
-    // Invalidate!
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
   ctx->extraAddressFixupInfo.intPtrOp = borrowedAddress.interiorPointerOp;
   auto borrowedValue = borrowedAddress.interiorPointerOp.getSingleBaseValue();
   if (!borrowedValue) {
-    // Invalidate!
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
@@ -929,16 +927,15 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
   // This cloner check must match the later cloner invocation in
   // replaceAddressUses()
   if (!canCloneUseDefChain(newValue, checkBase)) {
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
   // For now, just gather up uses
   auto &oldValueUses = ctx->extraAddressFixupInfo.allAddressUsesFromOldValue;
   if (InteriorPointerOperand::findTransitiveUsesForAddress(oldValue,
                                                            oldValueUses)) {
-    // If we found an error, invalidate and return!
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
@@ -1183,14 +1180,14 @@ OwnershipReplaceSingleUseHelper::OwnershipReplaceSingleUseHelper(
 
   // If we have an address, bail. We don't support this.
   if (newValue->getType().isAddress()) {
-    ctx = nullptr;
+    invalidate();
     return;
   }
 
   // Otherwise, lets check if we can perform this RAUW operation. If we can't,
   // set ctx to nullptr to invalidate the helper and return.
   if (!hasValidRAUWOwnership(use->get(), newValue)) {
-    ctx = nullptr;
+    invalidate();
     return;
   }
 

test/SILOptimizer/cse_ossa_nontrivial.sil

@@ -859,3 +859,37 @@ bb0(%0 : @guaranteed $TestKlass):
   return %139 : $()
 }
 
+struct X {
+  var i: Int64
+}
+
+sil_global hidden @globalX : $X
+
+sil [ossa] @use_address : $@convention(thin) (@in_guaranteed X) -> ()
+
+// CHECK-LABEL: sil [ossa] @cse_ossa_validrauwfollowinvalidrauw :
+// CHECK: global_addr
+// CHECK-NOT: global_addr
+// CHECK-LABEL: } // end sil function 'cse_ossa_validrauwfollowinvalidrauw'
+sil [ossa] @cse_ossa_validrauwfollowinvalidrauw : $@convention(thin) (@guaranteed TestKlass) -> () {
+bb0(%0 : @guaranteed $TestKlass):
+  %1 = function_ref @use_address : $@convention(thin) (@in_guaranteed X) -> ()
+  %2 = ref_element_addr %0 : $TestKlass, #TestKlass.testStruct
+  %3 = begin_access [modify] [dynamic] %2 : $*NonTrivialStruct
+  %4 = struct_element_addr %3 : $*NonTrivialStruct, #NonTrivialStruct.val
+  %5 = load_borrow %4 : $*Klass
+  %6 = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  %7 = apply %6(%5) : $@convention(thin) (@guaranteed Klass) -> ()
+  end_borrow %5 : $Klass
+  %9 = global_addr @globalX : $*X
+  %10 = apply %1(%9) : $@convention(thin) (@in_guaranteed X) -> ()
+  %11 = struct_element_addr %3 : $*NonTrivialStruct, #NonTrivialStruct.val
+  %12 = load_borrow %11 : $*Klass
+  %13 = apply %6(%12) : $@convention(thin) (@guaranteed Klass) -> ()
+  end_borrow %12 : $Klass
+  %15 = global_addr @globalX : $*X
+  %16 = apply %1(%15) : $@convention(thin) (@in_guaranteed X) -> ()
+  end_access %3 : $*NonTrivialStruct
+  %18 = tuple ()
+  return %18 : $()
+}