[RAUW] Allow RAUWing some lexical values.

RAUWing a lexical value with a non-lexical value is illegal because it
would result in the value's lifetime being shortened without regard to
deinit barriers.  RAUWing _with_ a lexical value is LEGAL so long as
doing so doesn't extend its lifetime.

Author: nate-chandler

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -212,7 +212,8 @@ class OwnershipRAUWHelper {
   /// 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);
+  static bool hasValidRAUWOwnership(SILValue oldValue, SILValue newValue,
+                                    ArrayRef<Operand *> oldUses);
 
 private:
   OwnershipFixupContext *ctx;
@@ -283,6 +284,10 @@ class OwnershipRAUWHelper {
   SILValue getReplacementAddress();
 };
 
+/// Whether the provided uses lie within the current liveness of the
+/// specified lexical value.
+bool areUsesWithinLexicalValueLifetime(SILValue, ArrayRef<Operand *>);
+
 /// A utility composed ontop of OwnershipFixupContext that knows how to replace
 /// a single use of a value with another value with a different ownership. We
 /// allow for the values to have different types.

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -297,13 +297,28 @@ static void cleanupOperandsBeforeDeletion(SILInstruction *oldValue,
 // directly checking ownership requirements. This does not determine whether the
 // scope of the newValue can be fully extended.
 bool OwnershipRAUWHelper::hasValidRAUWOwnership(SILValue oldValue,
-                                                SILValue newValue) {
+                                                SILValue newValue,
+                                                ArrayRef<Operand *> oldUses) {
   auto newOwnershipKind = newValue.getOwnershipKind();
 
   // If the either value is lexical, replacing its uses may result in
   // shortening or lengthening its lifetime in ways that don't respect lexical
   // scope and deinit barriers.
-  if (oldValue->isLexical() || newValue->isLexical())
+  //
+  // Specifically, we have the following cases:
+  //
+  // +--------+--------+
+  // |oldValue|newValue|
+  // +--------+--------+
+  // |  not   |  not   | legal
+  // +--------+--------+
+  // |lexical |  not   | illegal
+  // +--------+--------+
+  // |   *    |lexical | legal so long as it doesn't extend newValue's lifetime
+  // +--------+--------+
+  if ((oldValue->isLexical() && !newValue->isLexical()) ||
+      (newValue->isLexical() &&
+       !areUsesWithinLexicalValueLifetime(newValue, oldUses)))
     return false;
 
   // If our new kind is ValueOwnershipKind::None, then we are fine. We
@@ -353,20 +368,51 @@ bool OwnershipRAUWHelper::hasValidRAUWOwnership(SILValue oldValue,
 // extend the lifetime of \p oldValue to cover the new uses.
 static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue,
                                      OwnershipFixupContext &context) {
-  if (!OwnershipRAUWHelper::hasValidRAUWOwnership(oldValue, newValue))
-    return false;
+  switch (oldValue.getOwnershipKind()) {
+  case OwnershipKind::Guaranteed:
+    // Check that the old lifetime can be extended and record the necessary
+    // book-keeping in the OwnershipFixupContext.
+    context.clear();
+
+    // Check that no transitive uses have a PointerEscape, and record the leaf
+    // uses for liveness extension.
+    if (!findExtendedTransitiveGuaranteedUses(oldValue,
+                                              context.guaranteedUsePoints))
+      return false;
+    return OwnershipRAUWHelper::hasValidRAUWOwnership(
+        oldValue, newValue, context.guaranteedUsePoints);
+  default: {
+    SmallVector<Operand *, 8> ownedUsePoints;
+    // If newValue is lexical, find the uses of oldValue so that it can be
+    // determined whether the replacement would illegally extend the lifetime
+    // of newValue.
+    if (newValue->isLexical() &&
+        !findUsesOfSimpleValue(oldValue, &ownedUsePoints))
+      return false;
+    return OwnershipRAUWHelper::hasValidRAUWOwnership(oldValue, newValue,
+                                                      ownedUsePoints);
+  }
+  }
+}
+
+bool swift::areUsesWithinLexicalValueLifetime(SILValue value,
+                                              ArrayRef<Operand *> uses) {
+  assert(value->isLexical());
 
-  if (oldValue.getOwnershipKind() != OwnershipKind::Guaranteed)
+  // The lexical lifetime of a function argument is the whole body of the
+  // function.
+  if (isa<SILFunctionArgument>(value))
     return true;
 
-  // Check that the old lifetime can be extended and record the necessary
-  // book-keeping in the OwnershipFixupContext.
-  context.clear();
+  if (auto borrowedValue = BorrowedValue(value)) {
+    PrunedLiveness liveness;
+    auto *function = value->getFunction();
+    borrowedValue.computeLiveness(liveness);
+    DeadEndBlocks deadEndBlocks(function);
+    return liveness.areUsesWithinBoundary(uses, &deadEndBlocks);
+  }
 
-  // Check that no transitive uses have a PointerEscape, and record the leaf
-  // uses for liveness extension.
-  return findExtendedTransitiveGuaranteedUses(oldValue,
-                                              context.guaranteedUsePoints);
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1118,7 +1164,9 @@ SILValue OwnershipRAUWPrepare::prepareReplacement(SILValue newValue) {
   if (oldValue->use_empty())
     return newValue;
 
-  assert(OwnershipRAUWHelper::hasValidRAUWOwnership(oldValue, newValue) &&
+  assert(
+      OwnershipRAUWHelper::hasValidRAUWOwnership(oldValue, newValue,
+                                                 ctx.guaranteedUsePoints) &&
       "Should have checked if can perform this operation before calling it?!");
   // If our new value is just none, we can pass anything to it so just RAUW
   // and return.
@@ -1576,7 +1624,10 @@ OwnershipReplaceSingleUseHelper::OwnershipReplaceSingleUseHelper(
 
   // 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 (!OwnershipRAUWHelper::hasValidRAUWOwnership(use->get(), newValue)) {
+  SmallVector<Operand *, 1> oldUses;
+  oldUses.push_back(use);
+  if (!OwnershipRAUWHelper::hasValidRAUWOwnership(use->get(), newValue,
+                                                  oldUses)) {
     invalidate();
     return;
   }