Merge pull request swiftlang#35605 from meg-gupta/ossabugs1

Author: swift-ci

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;
             }
           }
@@ -239,8 +241,11 @@ static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue,
   if (oldValue.getOwnershipKind() == OwnershipKind::Guaranteed) {
     // Check that the old lifetime can be extended and record the necessary
     // book-keeping in the OwnershipFixupContext.
-    return findTransitiveBorrowedUses(oldValue, context.transitiveBorrowedUses,
-                                      context.recursiveReborrows);
+    if (!findTransitiveBorrowedUses(oldValue, context.transitiveBorrowedUses,
+                                    context.recursiveReborrows)) {
+      context.clear();
+      return false;
+    }
   }
   return true;
 }
@@ -430,18 +435,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 +508,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 +724,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,68 @@ 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 
+}
+
+// Test to make sure we clear the context if we fail while checking if ownership rauw is possible
+// CHECK-LABEL: sil [ossa] @test_rauwfailsandthensucceeds :
+// CHECK: bb0
+// CHECK: struct_extract
+// CHECK: struct_extract
+// CHECK: struct $_SliceBuffer
+// CHECK-NOT: struct $_SliceBuffer
+// CHECK-LABEL: } // end sil function 'test_rauwfailsandthensucceeds'
+sil [ossa] @test_rauwfailsandthensucceeds : $@convention(method) <Element> (UnsafeMutablePointer<Element>, @guaranteed _ContiguousArrayBuffer<Element>, Int, UInt) -> @owned _SliceBuffer<Element> {
+bb0(%0 : $UnsafeMutablePointer<Element>, %1 : @guaranteed $_ContiguousArrayBuffer<Element>, %2 : $Int, %3 : $UInt):
+  %4 = struct_extract %1 : $_ContiguousArrayBuffer<Element>, #_ContiguousArrayBuffer._storage
+  %5 = copy_value %4 : $__ContiguousArrayStorageBase
+  %6 = init_existential_ref %5 : $__ContiguousArrayStorageBase : $__ContiguousArrayStorageBase, $AnyObject
+  %7 = struct_extract %1 : $_ContiguousArrayBuffer<Element>, #_ContiguousArrayBuffer._storage
+  %8 = ref_tail_addr %7 : $__ContiguousArrayStorageBase, $Element
+  %9 = address_to_pointer %8 : $*Element to $Builtin.RawPointer
+  %10 = struct $UnsafeMutablePointer<Element> (%9 : $Builtin.RawPointer)
+  %11 = begin_borrow %6 : $AnyObject
+  %12 = struct $_SliceBuffer<Element> (%11 : $AnyObject, %0 : $UnsafeMutablePointer<Element>, %2 : $Int, %3 : $UInt)
+  %13 = copy_value %12 : $_SliceBuffer<Element>
+  end_borrow %11 : $AnyObject
+  destroy_value %13 : $_SliceBuffer<Element>
+  %14 = begin_borrow %6 : $AnyObject
+  %15 = struct $_SliceBuffer<Element> (%14 : $AnyObject, %0 : $UnsafeMutablePointer<Element>, %2 : $Int, %3 : $UInt)
+  %16 = copy_value %15 : $_SliceBuffer<Element>
+  end_borrow %14 : $AnyObject
+  %17 = struct $_SliceBuffer<Element> (%6 : $AnyObject, %0 : $UnsafeMutablePointer<Element>, %2 : $Int, %3 : $UInt)
+  destroy_value %17 : $_SliceBuffer<Element>
+  return %16 : $_SliceBuffer<Element>
+}
+