Fix use-after-free in ownership rauw

Instead of saving BorrowingOperand on the context save the SILBasicBlock
and index of the terminator operand.
This avoids the use-after-free in eliminateReborrowsOfRecursiveBorrows.
Previously, eliminateReborrowsOfRecursiveBorrows called helper
insertOwnedBaseValueAlongBranchEdge, which can delete a branch
instruction (reborrow) that could have been cached in
recursiveReborrows.

Author: meg-gupta

include/swift/SIL/OwnershipUtils.h

@@ -472,13 +472,14 @@ struct BorrowedValue {
   /// reborrows, place them in BorrowingOperand form into \p
   /// foundReborrows. Returns true if we appended any such reborrows to
   /// foundReborrows... false otherwise.
-  bool
-  gatherReborrows(SmallVectorImpl<BorrowingOperand> &foundReborrows) const {
+  bool gatherReborrows(SmallVectorImpl<std::pair<SILBasicBlock *, unsigned>>
+                           &foundReborrows) const {
     bool foundAnyReborrows = false;
     for (auto *op : value->getUses()) {
       if (auto borrowingOperand = BorrowingOperand::get(op)) {
         if (borrowingOperand.isReborrow()) {
-          foundReborrows.push_back(*borrowingOperand);
+          foundReborrows.push_back(
+              {value->getParentBlock(), op->getOperandNumber()});
           foundAnyReborrows = true;
         }
       }

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -39,7 +39,7 @@ struct OwnershipFixupContext {
   JointPostDominanceSetComputer &jointPostDomSetComputer;
 
   SmallVector<Operand *, 8> transitiveBorrowedUses;
-  SmallVector<BorrowingOperand, 8> recursiveReborrows;
+  SmallVector<std::pair<SILBasicBlock *, unsigned>, 8> recursiveReborrows;
 
   /// Extra state initialized by OwnershipRAUWFixupHelper::get() that we use
   /// when RAUWing addresses. This ensures we do not need to recompute this

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -100,8 +100,8 @@ insertOwnedBaseValueAlongBranchEdge(BranchInst *bi, SILValue innerCopy,
 }
 
 static bool findTransitiveBorrowedUses(
-  SILValue value, SmallVectorImpl<Operand *> &usePoints,
-  SmallVectorImpl<BorrowingOperand> &reborrowPoints) {
+    SILValue value, SmallVectorImpl<Operand *> &usePoints,
+    SmallVectorImpl<std::pair<SILBasicBlock *, unsigned>> &reborrowPoints) {
   assert(value.getOwnershipKind() == OwnershipKind::Guaranteed);
 
   unsigned firstOffset = usePoints.size();
@@ -165,7 +165,9 @@ static bool findTransitiveBorrowedUses(
         [&](Operand *scopeEndingUse) {
           if (auto scopeEndingBorrowingOp = BorrowingOperand(scopeEndingUse)) {
             if (scopeEndingBorrowingOp.isReborrow()) {
-              reborrowPoints.push_back(scopeEndingUse);
+              auto *branch = scopeEndingUse->getUser();
+              reborrowPoints.push_back(
+                  {branch->getParent(), scopeEndingUse->getOperandNumber()});
               return true;
             }
           }
@@ -430,18 +432,20 @@ OwnershipLifetimeExtender::createPlusZeroBorrow(SILValue newValue,
 //===----------------------------------------------------------------------===//
 
 static void eliminateReborrowsOfRecursiveBorrows(
-    ArrayRef<BorrowingOperand> transitiveReborrows,
+    ArrayRef<std::pair<SILBasicBlock *, unsigned>> transitiveReborrows,
     SmallVectorImpl<Operand *> &usePoints, InstModCallbacks &callbacks) {
   SmallVector<std::pair<SILPhiArgument *, SILPhiArgument *>, 8>
       baseBorrowedValuePair;
   // Ok, we have transitive reborrows.
-  for (auto borrowingOperand : transitiveReborrows) {
+  for (auto it : transitiveReborrows) {
     // We eliminate the reborrow by creating a new copy+borrow at the reborrow
     // edge from the base value and using that for the reborrow instead of the
     // actual value. We of course insert an end_borrow for our original incoming
     // value.
+    auto *bi = cast<BranchInst>(it.first->getTerminator());
+    auto &op = bi->getOperandRef(it.second);
+    BorrowingOperand borrowingOperand(&op);
     SILValue value = borrowingOperand->get();
-    auto *bi = cast<BranchInst>(borrowingOperand->getUser());
     SILBuilderWithScope reborrowBuilder(bi);
     // Use an auto-generated location here, because the branch may have an
     // incompatible LocationKind
@@ -501,15 +505,19 @@ static void eliminateReborrowsOfRecursiveBorrows(
   }
 }
 
-static void rewriteReborrows(SILValue newBorrowedValue,
-                             ArrayRef<BorrowingOperand> foundReborrows,
-                             InstModCallbacks &callbacks) {
+static void
+rewriteReborrows(SILValue newBorrowedValue,
+                 ArrayRef<std::pair<SILBasicBlock *, unsigned>> foundReborrows,
+                 InstModCallbacks &callbacks) {
   // Each initial reborrow that we have is a use of oldValue, so we know
   // that copy should be valid at the reborrow.
   SmallVector<std::pair<SILPhiArgument *, SILPhiArgument *>, 8>
       baseBorrowedValuePair;
-  for (auto reborrow : foundReborrows) {
-    auto *bi = cast<BranchInst>(reborrow.op->getUser());
+  for (auto it : foundReborrows) {
+    auto *bi = cast<BranchInst>(it.first->getTerminator());
+    auto &op = bi->getOperandRef(it.second);
+    BorrowingOperand reborrow(&op);
+
     SILBuilderWithScope reborrowBuilder(bi);
     // Use an auto-generated location here, because the branch may have an
     // incompatible LocationKind
@@ -713,7 +721,7 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleGuaranteed() {
   //    non-dominating copy value, allowing us to force our borrowing value to
   //    need a base phi argument (the one of our choosing).
   if (auto oldValueBorrowedVal = BorrowedValue::get(oldValue)) {
-    SmallVector<BorrowingOperand, 8> foundReborrows;
+    SmallVector<std::pair<SILBasicBlock *, unsigned>, 8> foundReborrows;
     if (oldValueBorrowedVal.gatherReborrows(foundReborrows)) {
       rewriteReborrows(newBorrowedValue, foundReborrows, ctx.callbacks);
     }

test/SILOptimizer/cse_ossa_nontrivial.sil

@@ -596,3 +596,37 @@ bb8:
   return %res : $()
 }
 
+struct RecurStruct {
+  var val:NonTrivialStruct
+}
+
+// Test to make sure we do not crash on a use after free when the branch instructions in bb1 and bb2 are deleted
+// CHECK-LABEL: sil [ossa] @test_useafterfreeinrauw :
+// CHECK: bb0
+// CHECK: struct_extract
+// CHECK-NOT: struct_extract
+// CHECK: cond_br
+// CHECK-LABEL: } // end sil function 'test_useafterfreeinrauw'
+sil [ossa] @test_useafterfreeinrauw : $@convention(thin) (@guaranteed RecurStruct) -> @owned NonTrivialStruct {
+bb0(%0 : @guaranteed $RecurStruct):
+  %1 = struct_extract %0 : $RecurStruct, #RecurStruct.val
+  debug_value %1 : $NonTrivialStruct
+  %2 = struct_extract %0 : $RecurStruct, #RecurStruct.val
+  cond_br undef, bb1, bb2
+
+bb1:
+  %3 = struct_extract %2 : $NonTrivialStruct, #NonTrivialStruct.val
+  %borrow3 = begin_borrow %3 : $Klass
+  br bb3(%borrow3 : $Klass)
+
+bb2:
+  %4 = struct_extract %2 : $NonTrivialStruct, #NonTrivialStruct.val
+  %borrow4 = begin_borrow %4 : $Klass
+  br bb3(%borrow4 : $Klass)
+
+bb3(%borrow : @guaranteed $Klass):
+  end_borrow %borrow : $Klass
+  %copy = copy_value %1 : $NonTrivialStruct
+  return %copy : $NonTrivialStruct 
+}
+