Merge pull request #39721 from atrick/oou-cleanup

OSSA RAUW cleanup in preparation for fixes and features

Author: atrick

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -647,13 +647,35 @@ FullApplySite cloneFullApplySiteReplacingCallee(FullApplySite applySite,
                                                 SILValue newCallee,
                                                 SILBuilderContext &builderCtx);
 
+/// Replace all uses of \p oldValue with \p newValue, notifying the callbacks
+/// of new uses and when end-of-scope instructions are deleted.
+SILBasicBlock::iterator replaceAllUses(SILValue oldValue, SILValue newValue,
+                                       SILBasicBlock::iterator nextii,
+                                       InstModCallbacks &callbacks);
+
 /// This is a low level routine that makes all uses of \p svi uses of \p
 /// newValue (ignoring end scope markers) and then deletes \p svi and all end
 /// scope markers. Then returns the next inst to process.
 SILBasicBlock::iterator replaceAllUsesAndErase(SingleValueInstruction *svi,
                                                SILValue newValue,
                                                InstModCallbacks &callbacks);
 
+/// Replace all uses of \p oldValue with \p newValue, delete the instruction
+/// that defines \p oldValue, and notify the callbacks of new uses and when
+/// the defining instruction and its end-of-scope instructions are deleted.
+///
+/// Precondition: \p oldValue must be a SingleValueInstruction or a terminator
+/// result. \p oldValue must be the only result with remaining uses. For
+/// terminators with multiple results, remove all other results for, e.g. via
+/// replaceAllUsesWithUndef().
+///
+/// If \p oldValue is a terminator result, a new branch instruction is inserted
+/// in place of the old terminator and all basic block successors become
+/// unreachable except for the successor containing the replaced result.
+SILBasicBlock::iterator replaceAllUsesAndErase(SILValue oldValue,
+                                               SILValue newValue,
+                                               InstModCallbacks &callbacks);
+
 /// This API is equivalent to performing \p use->set(\p newValue) except that:
 ///
 /// 1. If the user of \p use is an end scope, this API no-opts. This API is only

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -104,29 +104,46 @@ struct OwnershipFixupContext {
   }
 };
 
-/// A utility composed ontop of OwnershipFixupContext that knows how to RAUW a
+/// A utility composed on top of OwnershipFixupContext that knows how to RAUW a
 /// value or a single value instruction with a new value and then fixup
 /// ownership invariants afterwards.
 class OwnershipRAUWHelper {
+public:
+  /// Return true if \p oldValue can be replaced with \p newValue in terms of
+  /// their value ownership. This ignores any current uses of \p oldValue. To
+  /// determine whether \p oldValue can be replaced as-is with it's existing
+  /// uses, create an instance of OwnershipRAUWHelper and check its validity.
+  static bool hasValidRAUWOwnership(SILValue oldValue, SILValue newValue);
+
+private:
   OwnershipFixupContext *ctx;
-  SingleValueInstruction *oldValue;
+  SILValue oldValue;
   SILValue newValue;
 
 public:
-  OwnershipRAUWHelper() : ctx(nullptr), oldValue(nullptr), newValue(nullptr) {}
+  OwnershipRAUWHelper() : ctx(nullptr) {}
+
+  ~OwnershipRAUWHelper() { if (ctx) ctx->clear(); }
 
   /// Return an instance of this class if we can perform the specific RAUW
   /// operation ignoring if the types line up. Returns None otherwise.
   ///
+  /// \p oldValue may be either a SingleValueInstruction or a terminator result.
+  ///
+  /// Precondition: If \p oldValue is a BorrowedValue that introduces a local
+  /// borrow scope, then \p newValue must either be defined in the same block as
+  /// \p oldValue, or it must dominate \p oldValue (rather than merely
+  /// dominating its uses).
+  ///
   /// DISCUSSION: We do not check that the types line up here so that we can
   /// allow for our users to transform our new value in ways that preserve
   /// ownership at \p oldValue before we perform the actual RAUW. If \p newValue
   /// is an object, any instructions in the chain of transforming instructions
   /// from \p newValue at \p oldValue's must be forwarding. If \p newValue is an
   /// address, then these transforms can only transform the address into a
   /// derived address.
-  OwnershipRAUWHelper(OwnershipFixupContext &ctx,
-                      SingleValueInstruction *oldValue, SILValue newValue);
+  OwnershipRAUWHelper(OwnershipFixupContext &ctx, SILValue oldValue,
+                      SILValue newValue);
 
   /// Returns true if this helper was initialized into a valid state.
   operator bool() const { return isValid(); }
@@ -149,7 +166,6 @@ class OwnershipRAUWHelper {
                                              SILValue newValue);
 
   void invalidate() {
-    ctx->clear();
     ctx = nullptr;
   }
 };
@@ -158,6 +174,9 @@ class OwnershipRAUWHelper {
 /// a single use of a value with another value with a different ownership. We
 /// allow for the values to have different types.
 ///
+/// Precondition: if \p use ends a borrow scope, then \p newValue dominates the
+/// BorrowedValue that begins the scope.
+///
 /// NOTE: When not in OSSA, this just performs a normal set use, so this code is
 /// safe to use with all code.
 class OwnershipReplaceSingleUseHelper {
@@ -177,6 +196,8 @@ class OwnershipReplaceSingleUseHelper {
   OwnershipReplaceSingleUseHelper(OwnershipFixupContext &ctx, Operand *use,
                                   SILValue newValue);
 
+  ~OwnershipReplaceSingleUseHelper() { if (ctx) ctx->clear(); }
+
   /// Returns true if this helper was initialized into a valid state.
   operator bool() const { return isValid(); }
   bool isValid() const { return bool(ctx) && bool(use) && bool(newValue); }

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -4018,6 +4018,8 @@ bool SimplifyCFG::simplifyArgument(SILBasicBlock *BB, unsigned i) {
   return true;
 }
 
+// OWNERSHIP NOTE: This is always safe for guaranteed and owned arguments since
+// in both cases the phi will consume its input.
 static void tryToReplaceArgWithIncomingValue(SILBasicBlock *BB, unsigned i,
                                              DominanceInfo *DT) {
   auto *A = BB->getArgument(i);

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1925,15 +1925,19 @@ swift::cloneFullApplySiteReplacingCallee(FullApplySite applySite,
   llvm_unreachable("Unhandled case?!");
 }
 
-SILBasicBlock::iterator
-swift::replaceAllUsesAndErase(SingleValueInstruction *svi, SILValue newValue,
-                              InstModCallbacks &callbacks) {
-  assert(svi != newValue && "Cannot RAUW a value with itself");
-  SILBasicBlock::iterator nextii = std::next(svi->getIterator());
-
-  // Only SingleValueInstructions are currently simplified.
-  while (!svi->use_empty()) {
-    Operand *use = *svi->use_begin();
+// FIXME: For any situation where this may be called on an unbounded number of
+// uses, it should perform a single callback invocation to notify the client
+// that newValue has new uses rather than a callback for every new use.
+//
+// FIXME: This should almost certainly replace end_lifetime uses rather than
+// deleting them.
+SILBasicBlock::iterator swift::replaceAllUses(SILValue oldValue,
+                                              SILValue newValue,
+                                              SILBasicBlock::iterator nextii,
+                                              InstModCallbacks &callbacks) {
+  assert(oldValue != newValue && "Cannot RAUW a value with itself");
+  while (!oldValue->use_empty()) {
+    Operand *use = *oldValue->use_begin();
     SILInstruction *user = use->getUser();
     // Erase the end of scope marker.
     if (isEndOfScopeMarker(user)) {
@@ -1944,12 +1948,62 @@ swift::replaceAllUsesAndErase(SingleValueInstruction *svi, SILValue newValue,
     }
     callbacks.setUseValue(use, newValue);
   }
+  return nextii;
+}
+
+SILBasicBlock::iterator
+swift::replaceAllUsesAndErase(SingleValueInstruction *svi, SILValue newValue,
+                              InstModCallbacks &callbacks) {
+  SILBasicBlock::iterator nextii = replaceAllUses(
+      svi, newValue, std::next(svi->getIterator()), callbacks);
 
   callbacks.deleteInst(svi);
 
   return nextii;
 }
 
+SILBasicBlock::iterator
+swift::replaceAllUsesAndErase(SILValue oldValue, SILValue newValue,
+                              InstModCallbacks &callbacks) {
+  auto *blockArg = dyn_cast<SILPhiArgument>(oldValue);
+  if (!blockArg) {
+    // SingleValueInstruction SSA replacement.
+    return replaceAllUsesAndErase(cast<SingleValueInstruction>(oldValue),
+                                  newValue, callbacks);
+  }
+  llvm_unreachable("Untested");
+#if 0 // FIXME: to be enabled in a following commit
+  TermInst *oldTerm = blockArg->getTerminatorForResult();
+  assert(oldTerm && "can only replace and erase terminators, not phis");
+
+  // Before:
+  //     oldTerm bb1, bb2
+  //   bb1(%oldValue):
+  //     use %oldValue
+  //   bb2:
+  //
+  // After:
+  //     br bb1
+  //   bb1:
+  //     use %newValue
+  //   bb2:
+
+  auto nextii = replaceAllUses(blockArg, newValue,
+                               oldTerm->getParent()->end(), callbacks);
+  // Now that oldValue is replaced, the terminator should have no uses
+  // left. The caller should have removed uses from other results.
+  for (auto *succBB : oldTerm->getSuccessorBlocks()) {
+    assert(succBB->getNumArguments() == 1 && "expected terminator result");
+    succBB->eraseArgument(0);
+  }
+  auto *newBr = SILBuilderWithScope(oldTerm).createBranch(
+    oldTerm->getLoc(), blockArg->getParent());
+  callbacks.createdNewInst(newBr);
+  callbacks.deleteInst(oldTerm);
+  return nextii;
+#endif
+}
+
 /// Given that we are going to replace use's underlying value, if the use is a
 /// lifetime ending use, insert an end scope scope use for the underlying value
 /// before we RAUW.