Fix ownership rauw to not leave behind stale ownership fixup context

Ownership rauw uses a shared ownership fixup context to maintain state.
When ownership rauw fails, due to some invalid condition, we leave
behind stale data in this shared ownership fixup context.
This stale context can indvertantly affect the next rauw on addresses.
In addition to setting the ownership fixup context to nullptr, we
should also clear it so that it's internal data structures are
cleared.

Author: meg-gupta

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 : $()
+}