[sil-combine] Fix optimization of COWBufferRead for ownership

include/swift/SIL/OwnershipUtils.h

@@ -527,13 +527,34 @@ struct BorrowedValue {
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 
   /// Visit each of the interior pointer uses of this underlying borrow
-  /// introduced value. These object -> address projections and any transitive
-  /// address uses must be treated as liveness requiring uses of the guaranteed
-  /// value and we can not shrink the scope beyond that point. Returns true if
-  /// we were able to understand all uses and thus guarantee we found all
-  /// interior pointer uses. Returns false otherwise.
+  /// introduced value without looking through nested borrows or reborrows.
+  ///
+  /// These object -> address projections and any transitive address uses must
+  /// be treated as liveness requiring uses of the guaranteed value and we can
+  /// not shrink the scope beyond that point. Returns true if we were able to
+  /// understand all uses and thus guarantee we found all interior pointer
+  /// uses. Returns false otherwise.
   bool visitInteriorPointerOperands(
-      function_ref<void(const InteriorPointerOperand &)> func) const;
+      function_ref<void(InteriorPointerOperand)> func) const {
+    return visitInteriorPointerOperandHelper(
+        func, InteriorPointerOperandVisitorKind::NoNestedNoReborrows);
+  }
+
+  /// Visit each of the interior pointer uses of this underlying borrow
+  /// introduced value looking through nested borrow scopes but not reborrows.
+  bool visitNestedInteriorPointerOperands(
+      function_ref<void(InteriorPointerOperand)> func) const {
+    return visitInteriorPointerOperandHelper(
+        func, InteriorPointerOperandVisitorKind::YesNestedNoReborrows);
+  }
+
+  /// Visit each of the interior pointer uses of this underlying borrow
+  /// introduced value looking through nested borrow scopes and reborrows.
+  bool visitExtendedInteriorPointerOperands(
+      function_ref<void(InteriorPointerOperand)> func) const {
+    return visitInteriorPointerOperandHelper(
+        func, InteriorPointerOperandVisitorKind::YesNestedYesReborrows);
+  }
 
   /// Visit all immediate uses of this borrowed value and if any of them are
   /// reborrows, place them in BorrowingOperand form into \p
@@ -560,6 +581,16 @@ struct BorrowedValue {
   SILValue operator->() const { return value; }
   SILValue operator*() { return value; }
   SILValue operator*() const { return value; }
+
+private:
+  enum class InteriorPointerOperandVisitorKind {
+    NoNestedNoReborrows,
+    YesNestedNoReborrows,
+    YesNestedYesReborrows,
+  };
+  bool visitInteriorPointerOperandHelper(
+      function_ref<void(InteriorPointerOperand)> func,
+      InteriorPointerOperandVisitorKind kind) const;
 };
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
@@ -603,6 +634,8 @@ class InteriorPointerOperandKind {
     return value;
   }
 
+  explicit operator bool() const { return isValid(); }
+
   bool isValid() const { return value != Kind::Invalid; }
 
   static InteriorPointerOperandKind get(Operand *use) {
@@ -651,7 +684,6 @@ struct InteriorPointerOperand {
 
   InteriorPointerOperand(Operand *op)
       : operand(op), kind(InteriorPointerOperandKind::get(op)) {
-    assert(kind.isValid());
   }
 
   operator bool() const {

lib/SIL/Utils/OwnershipUtils.cpp

@@ -600,24 +600,56 @@ bool BorrowedValue::visitExtendedLocalScopeEndingUses(
   return true;
 }
 
-bool BorrowedValue::visitInteriorPointerOperands(
-    function_ref<void(const InteriorPointerOperand &)> func) const {
+bool BorrowedValue::visitInteriorPointerOperandHelper(
+    function_ref<void(InteriorPointerOperand)> func,
+    BorrowedValue::InteriorPointerOperandVisitorKind kind) const {
+  using Kind = BorrowedValue::InteriorPointerOperandVisitorKind;
+
   SmallVector<Operand *, 32> worklist(value->getUses());
   while (!worklist.empty()) {
     auto *op = worklist.pop_back_val();
 
-    if (auto interiorPointer = InteriorPointerOperand::get(op)) {
+    if (auto interiorPointer = InteriorPointerOperand(op)) {
       func(interiorPointer);
       continue;
     }
 
+    if (auto borrowingOperand = BorrowingOperand(op)) {
+      switch (kind) {
+      case Kind::NoNestedNoReborrows:
+        // We do not look through nested things and or reborrows, so just
+        // continue.
+        continue;
+      case Kind::YesNestedNoReborrows:
+        // We only look through nested borrowing operands, we never look through
+        // reborrows though.
+        if (borrowingOperand.isReborrow())
+          continue;
+        break;
+      case Kind::YesNestedYesReborrows:
+        // Look through everything!
+        break;
+      }
+
+      borrowingOperand.visitBorrowIntroducingUserResults([&](auto bv) {
+        for (auto *use : bv->getUses()) {
+          if (auto intPtrOperand = InteriorPointerOperand(use)) {
+            func(intPtrOperand);
+            continue;
+          }
+          worklist.push_back(use);
+        }
+        return true;
+      });
+      continue;
+    }
+
     auto *user = op->getUser();
-    if (isa<BeginBorrowInst>(user) || isa<DebugValueInst>(user) ||
-        isa<SuperMethodInst>(user) || isa<ClassMethodInst>(user) ||
-        isa<CopyValueInst>(user) || isa<EndBorrowInst>(user) ||
-        isa<ApplyInst>(user) || isa<StoreBorrowInst>(user) ||
-        isa<StoreInst>(user) || isa<PartialApplyInst>(user) ||
-        isa<UnmanagedRetainValueInst>(user) ||
+    if (isa<DebugValueInst>(user) || isa<SuperMethodInst>(user) ||
+        isa<ClassMethodInst>(user) || isa<CopyValueInst>(user) ||
+        isa<EndBorrowInst>(user) || isa<ApplyInst>(user) ||
+        isa<StoreBorrowInst>(user) || isa<StoreInst>(user) ||
+        isa<PartialApplyInst>(user) || isa<UnmanagedRetainValueInst>(user) ||
         isa<UnmanagedReleaseValueInst>(user) ||
         isa<UnmanagedAutoreleaseValueInst>(user)) {
       continue;
@@ -693,10 +725,11 @@ bool InteriorPointerOperand::findTransitiveUsesForAddress(
     if (Projection::isAddressProjection(user) ||
         isa<ProjectBlockStorageInst>(user) ||
         isa<OpenExistentialAddrInst>(user) ||
-        isa<InitExistentialAddrInst>(user) ||
-        isa<InitEnumDataAddrInst>(user) || isa<BeginAccessInst>(user) ||
-        isa<TailAddrInst>(user) || isa<IndexAddrInst>(user) ||
-        isa<UnconditionalCheckedCastAddrInst>(user)) {
+        isa<InitExistentialAddrInst>(user) || isa<InitEnumDataAddrInst>(user) ||
+        isa<BeginAccessInst>(user) || isa<TailAddrInst>(user) ||
+        isa<IndexAddrInst>(user) ||
+        isa<UnconditionalCheckedCastAddrInst>(user) ||
+        isa<UncheckedAddrCastInst>(user)) {
       for (SILValue r : user->getResults()) {
         llvm::copy(r->getUses(), std::back_inserter(worklist));
       }

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -312,6 +312,8 @@ class SILCombiner :
   SILInstruction *optimizeBuiltinIsConcrete(BuiltinInst *I);
 
   SILInstruction *optimizeBuiltinCOWBufferForReading(BuiltinInst *BI);
+  SILInstruction *optimizeBuiltinCOWBufferForReadingNonOSSA(BuiltinInst *BI);
+  SILInstruction *optimizeBuiltinCOWBufferForReadingOSSA(BuiltinInst *BI);
 
   // Optimize the "trunc_N1_M2" builtin. if N1 is a result of "zext_M1_*" and
   // the following holds true: N1 > M1 and M2>= M1

lib/SILOptimizer/SILCombiner/SILCombinerBuiltinVisitors.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "sil-combine"
+
 #include "SILCombiner.h"
 #include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/DynamicCasts.h"
@@ -22,6 +23,7 @@
 #include "swift/SILOptimizer/Analysis/ValueTracking.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -109,6 +111,92 @@ SILInstruction *SILCombiner::optimizeBuiltinIsConcrete(BuiltinInst *BI) {
   return Builder.createIntegerLiteral(BI->getLoc(), BI->getType(), 1);
 }
 
+/// Replace
+/// \code
+///   %b = builtin "COWBufferForReading" %r
+///   %bb = begin_borrow %b
+///   %a = ref_element_addr %bb
+///   ... use %a ...
+///   end_borrow %bb
+/// \endcode
+/// with
+/// \code
+///   %bb = begin_borrow %r
+///   %a = ref_element_addr [immutable] %r
+///   ... use %b ...
+///   end_borrow %bb
+/// \endcode
+/// The same for ref_tail_addr.
+SILInstruction *
+SILCombiner::optimizeBuiltinCOWBufferForReadingOSSA(BuiltinInst *bi) {
+  SmallVector<BorrowedValue, 32> accumulatedBorrowedValues;
+
+  // A helper that performs our main loop to look through uses. It ensures
+  // that we do not need to fill up the useWorklist on the first iteration.
+  for (auto *use : bi->getUses()) {
+    // See if we have a borrowing operand that we can find a local borrowed
+    // value for. In such a case, we stash that borrowed value so that we can
+    // use it to find interior pointer operands.
+    if (auto operand = BorrowingOperand(use)) {
+      if (operand.isReborrow())
+        return nullptr;
+      operand.visitBorrowIntroducingUserResults([&](BorrowedValue bv) {
+        accumulatedBorrowedValues.push_back(bv);
+        return true;
+      });
+      continue;
+    }
+
+    // Otherwise, look for instructions that we know are uses that we can
+    // ignore.
+    auto *user = use->getUser();
+
+    // Debug instructions are safe.
+    if (user->isDebugInstruction())
+      continue;
+
+    // copy_value, destroy_value are safe due to our checking of the
+    // instruction use list for safety.
+    if (isa<DestroyValueInst>(user) || isa<CopyValueInst>(user))
+      continue;
+
+    // An instruction we don't understand, bail.
+    return nullptr;
+  }
+
+  // Now that we know that we have a case we support, use our stashed
+  // BorrowedValues to find all interior pointer operands into this copy of our
+  // COWBuffer and mark them as immutable.
+  //
+  // NOTE: We currently only use nested int ptr operands instead of extended int
+  // ptr operands since we do not want to look through reborrows and thus lose
+  // dominance.
+  while (!accumulatedBorrowedValues.empty()) {
+    auto bv = accumulatedBorrowedValues.pop_back_val();
+    bv.visitNestedInteriorPointerOperands(
+        [&](InteriorPointerOperand intPtrOperand) {
+          switch (intPtrOperand.kind) {
+          case InteriorPointerOperandKind::Invalid:
+            llvm_unreachable("Invalid int pointer kind?!");
+          case InteriorPointerOperandKind::RefElementAddr:
+            cast<RefElementAddrInst>(intPtrOperand->getUser())->setImmutable();
+            return;
+          case InteriorPointerOperandKind::RefTailAddr:
+            cast<RefTailAddrInst>(intPtrOperand->getUser())->setImmutable();
+            return;
+          case InteriorPointerOperandKind::OpenExistentialBox:
+            // Can not mark this immutable.
+            return;
+          }
+        });
+  }
+
+  OwnershipRAUWHelper helper(ownershipFixupContext, bi, bi->getOperand(0));
+  assert(helper && "COWBufferForReading always has an owned arg/owned result");
+  helper.perform();
+  return nullptr;
+}
+
 /// Replace
 /// \code
 ///   %b = builtin "COWBufferForReading" %r
@@ -119,23 +207,24 @@ SILInstruction *SILCombiner::optimizeBuiltinIsConcrete(BuiltinInst *BI) {
 ///   %a = ref_element_addr [immutable] %r
 /// \endcode
 /// The same for ref_tail_addr.
-SILInstruction *SILCombiner::optimizeBuiltinCOWBufferForReading(BuiltinInst *BI) {
-  auto useIter = BI->use_begin();
-  while (useIter != BI->use_end()) {
+SILInstruction *
+SILCombiner::optimizeBuiltinCOWBufferForReadingNonOSSA(BuiltinInst *bi) {
+  auto useIter = bi->use_begin();
+  while (useIter != bi->use_end()) {
     auto nextIter = std::next(useIter);
     SILInstruction *user = useIter->getUser();
-    SILValue ref = BI->getOperand(0);
+    SILValue ref = bi->getOperand(0);
     switch (user->getKind()) {
       case SILInstructionKind::RefElementAddrInst: {
-        auto *REAI = cast<RefElementAddrInst>(user);
-        REAI->setOperand(ref);
-        REAI->setImmutable();
+        auto *reai = cast<RefElementAddrInst>(user);
+        reai->setOperand(ref);
+        reai->setImmutable();
         break;
       }
       case SILInstructionKind::RefTailAddrInst: {
-        auto *RTAI = cast<RefTailAddrInst>(user);
-        RTAI->setOperand(ref);
-        RTAI->setImmutable();
+        auto *rtai = cast<RefTailAddrInst>(user);
+        rtai->setOperand(ref);
+        rtai->setImmutable();
         break;
       }
     case SILInstructionKind::DestroyValueInst:
@@ -151,11 +240,18 @@ SILInstruction *SILCombiner::optimizeBuiltinCOWBufferForReading(BuiltinInst *BI)
   }
 
   // If there are unknown users, keep the builtin, and IRGen will handle it.
-  if (BI->use_empty())
-    return eraseInstFromFunction(*BI);
+  if (bi->use_empty())
+    return eraseInstFromFunction(*bi);
   return nullptr;
 }
 
+SILInstruction *
+SILCombiner::optimizeBuiltinCOWBufferForReading(BuiltinInst *BI) {
+  if (hasOwnership())
+    return optimizeBuiltinCOWBufferForReadingOSSA(BI);
+  return optimizeBuiltinCOWBufferForReadingNonOSSA(BI);
+}
+
 static unsigned getTypeWidth(SILType Ty) {
   if (auto BuiltinIntTy = Ty.getAs<BuiltinIntegerType>()) {
     if (BuiltinIntTy->isFixedWidth()) {