[ownership] Add a new API to OwnershipFixupContext::replaceAllAddressUsesFixingInteriorPointerOwnership.

In OSSA, we enforce that addresses from interior pointer instructions are scoped
within a borrow scope. This means that it is invalid to use such an address
outside of its parent borrow scope and as a result one can not just RAUW an
address value by a dominating address value since the latter may be invalid at
the former. I foresee that I am going to have to solve this problem and so I
decided to write this API to handle the vast majority of cases.

The way this API works is that it:

1. Computes an access path with base for the new value. If we do not have a base
value and a valid access path with root, we bail.

2. Then we check if our base value is the result of an interior pointer
instruction. If it isn't, we are immediately done and can RAUW without further
delay.

3. If we do have an interior pointer instruction, we see if the immediate
guaranteed value we projected from has a single borrow introducer value. If not,
we bail. I think this is reasonable since with time, all guaranteed values will
always only have a single borrow introducing value (once struct, tuple,
destructure_struct, destructure_tuple become reborrows).

4. Then we gather up all inner uses of our access path. If for some reason that
fails, we bail.

5. Then we see if all of those uses are within our borrow scope. If so, we can
RAUW without any further worry.

6. Otherwise, we perform a copy+borrow of our interior pointer's operand value
at the interior pointer, create a copy of the interior pointer instruction upon
this new borrow and then RAUW oldValue with that instead. By construction all
uses of oldValue will be within this new interior pointer scope.

Author: gottesmm

include/swift/SIL/OwnershipUtils.h

@@ -669,7 +669,17 @@ struct InteriorPointerOperand {
   /// requirements to ensure that the underlying class is alive at all use
   /// points.
   bool getImplicitUses(SmallVectorImpl<Operand *> &foundUses,
-                       std::function<void(Operand *)> *onError = nullptr);
+                       std::function<void(Operand *)> *onError = nullptr) {
+    return getImplicitUsesForAddress(getProjectedAddress(), foundUses, onError);
+  }
+
+  /// The algorithm that is used to determine what the verifier will consider to
+  /// be implicit uses of the given address. Used to implement \see
+  /// getImplicitUses.
+  static bool
+  getImplicitUsesForAddress(SILValue address,
+                            SmallVectorImpl<Operand *> &foundUses,
+                            std::function<void(Operand *)> *onError = nullptr);
 
   Operand *operator->() { return operand; }
   const Operand *operator->() const { return operand; }

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -38,39 +38,89 @@ struct OwnershipFixupContext {
   DeadEndBlocks &deBlocks;
   JointPostDominanceSetComputer &jointPostDomSetComputer;
 
-  OwnershipFixupContext(InstModCallbacks &callbacks, DeadEndBlocks &deBlocks,
-                        JointPostDominanceSetComputer &inputJPDComputer)
-      : inlineCallbacks(), callbacks(callbacks), deBlocks(deBlocks),
-        jointPostDomSetComputer(inputJPDComputer) {}
+  /// Extra state initialized by OwnershipRAUWFixupHelper::get() that we use
+  /// when RAUWing addresses. This ensures we do not need to recompute this
+  /// state when we perform the actual RAUW.
+  struct AddressFixupContext {
+    /// When determining if we need to perform an address pointer fixup, we
+    /// compute all transitive address uses of oldValue. If we find that we do
+    /// need this fixed up, then we will copy our interior pointer base value
+    /// and use this to seed that new lifetime.
+    SmallVector<Operand *, 8> allAddressUsesFromOldValue;
+
+    /// This is the interior pointer operand that the new value we want to RAUW
+    /// is transitively derived from and enables us to know the underlying
+    /// borrowed base value that we need to lifetime extend.
+    InteriorPointerOperand intPtrOp;
 
-  OwnershipFixupContext(DeadEndBlocks &deBlocks,
-                        JointPostDominanceSetComputer &inputJPDComputer)
-      : inlineCallbacks(InstModCallbacks()), callbacks(*inlineCallbacks),
-        deBlocks(deBlocks), jointPostDomSetComputer(inputJPDComputer) {}
+    void clear() {
+      allAddressUsesFromOldValue.clear();
+      intPtrOp = InteriorPointerOperand();
+    }
+  };
+  AddressFixupContext extraAddressFixupInfo;
+
+  OwnershipFixupContext(InstModCallbacks &callbacks, DeadEndBlocks &deBlocks,
+                        JointPostDominanceSetComputer &jointPostDomSetComputer)
+      : callbacks(callbacks), deBlocks(deBlocks),
+        jointPostDomSetComputer(jointPostDomSetComputer) {}
 
   void clear() {
     jointPostDomSetComputer.clear();
+    extraAddressFixupInfo.allAddressUsesFromOldValue.clear();
+    extraAddressFixupInfo.intPtrOp = InteriorPointerOperand();
+  }
+
+private:
+  /// Helper method called to determine if we discovered we needed interior
+  /// pointer fixups while simplifying.
+  bool needsInteriorPointerFixups() const {
+    return bool(extraAddressFixupInfo.intPtrOp);
   }
 };
 
 /// A utility composed ontop 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 {
-  OwnershipFixupContext &ctx;
+  OwnershipFixupContext *ctx;
+  SingleValueInstruction *oldValue;
+  SILValue newValue;
 
 public:
-  OwnershipRAUWHelper(OwnershipFixupContext &ctx) : ctx(ctx) {}
+  OwnershipRAUWHelper() : ctx(nullptr), oldValue(nullptr), newValue(nullptr) {}
+
+  /// Return an instance of this class if we can perform the specific RAUW
+  /// operation ignoring if the types line up. Returns None otherwise.
+  ///
+  /// 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);
 
-  SILBasicBlock::iterator
-  replaceAllUsesAndErase(SingleValueInstruction *oldValue, SILValue newValue);
+  /// Returns true if this helper was initialized into a valid state.
+  operator bool() const { return isValid(); }
+  bool isValid() const { return bool(ctx) && bool(oldValue) && bool(newValue); }
 
-  /// We can not RAUW all old values with new values.
+  /// Perform the actual RAUW. We require that \p newValue and \p oldValue have
+  /// the same type at this point (in contrast to when calling
+  /// OwnershipRAUWFixupHelper::get()).
   ///
-  /// Namely, we do not support RAUWing values with ValueOwnershipKind::None
-  /// that have uses that do not require ValueOwnershipKind::None or
-  /// ValueOwnershipKind::Any.
-  static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue);
+  /// This is so that we can avoid creating "forwarding" transformation
+  /// instructions before we know if we can perform the RAUW. Any such
+  /// "forwarding" transformation must be performed upon \p newValue at \p
+  /// oldValue's insertion point so that we can then here RAUW the transformed
+  /// \p newValue.
+  SILBasicBlock::iterator perform();
+
+private:
+  SILBasicBlock::iterator replaceAddressUses(SingleValueInstruction *oldValue,
+                                             SILValue newValue);
 };
 
 } // namespace swift

lib/SIL/Utils/OwnershipUtils.cpp

@@ -520,10 +520,9 @@ bool BorrowedValue::visitInteriorPointerOperands(
 //                           InteriorPointerOperand
 //===----------------------------------------------------------------------===//
 
-bool InteriorPointerOperand::getImplicitUses(
-    SmallVectorImpl<Operand *> &foundUses,
+bool InteriorPointerOperand::getImplicitUsesForAddress(
+    SILValue projectedAddress, SmallVectorImpl<Operand *> &foundUses,
     std::function<void(Operand *)> *onError) {
-  SILValue projectedAddress = getProjectedAddress();
   SmallVector<Operand *, 8> worklist(projectedAddress->getUses());
 
   bool foundError = false;

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -755,8 +755,8 @@ swift::replaceAllSimplifiedUsesAndErase(SILInstruction *i, SILValue result,
   if (svi->getFunction()->hasOwnership()) {
     JointPostDominanceSetComputer computer(*deadEndBlocks);
     OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
-    OwnershipRAUWHelper helper(ctx);
-    return helper.replaceAllUsesAndErase(svi, result);
+    OwnershipRAUWHelper helper(ctx, svi, result);
+    return helper.perform();
   }
   return replaceAllUsesAndErase(svi, result, callbacks);
 }
@@ -781,20 +781,7 @@ SILValue swift::simplifyOverflowBuiltinInstruction(BuiltinInst *BI) {
 /// NOTE: We assume that the insertion point associated with the SILValue must
 /// dominate \p i.
 static SILValue simplifyInstruction(SILInstruction *i) {
-  SILValue result = InstSimplifier().visit(i);
-  if (!result)
-    return SILValue();
-
-  // If we have a result, we know that we must have a single value instruction
-  // by assumption since we have not implemented support in the rest of inst
-  // simplify for non-single value instructions. We put the cast here so that
-  // this code is not updated at this point in time.
-  auto *svi = cast<SingleValueInstruction>(i);
-  if (svi->getFunction()->hasOwnership())
-    if (!OwnershipRAUWHelper::canFixUpOwnershipForRAUW(svi, result))
-      return SILValue();
-
-  return result;
+  return InstSimplifier().visit(i);
 }
 
 SILBasicBlock::iterator swift::simplifyAndReplaceAllSimplifiedUsesAndErase(
@@ -811,11 +798,16 @@ SILBasicBlock::iterator swift::simplifyAndReplaceAllSimplifiedUsesAndErase(
   if (!result || svi == result)
     return next;
 
-  if (svi->getFunction()->hasOwnership()) {
-    JointPostDominanceSetComputer computer(*deadEndBlocks);
-    OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
-    OwnershipRAUWHelper helper(ctx);
-    return helper.replaceAllUsesAndErase(svi, result);
-  }
-  return replaceAllUsesAndErase(svi, result, callbacks);
+  if (!svi->getFunction()->hasOwnership())
+    return replaceAllUsesAndErase(svi, result, callbacks);
+
+  JointPostDominanceSetComputer computer(*deadEndBlocks);
+  OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
+  OwnershipRAUWHelper helper(ctx, svi, result);
+
+  // If our RAUW helper is invalid, we do not support RAUWing this case, so
+  // just return next.
+  if (!helper.isValid())
+    return next;
+  return helper.perform();
 }

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -96,11 +96,8 @@ class SILCombiner :
   /// post-dominating blocks to a full jointly post-dominating set.
   JointPostDominanceSetComputer jPostDomComputer;
 
-  /// A utility that we use to perform erase+RAUW that fixes up ownership for us
-  /// afterwards by lifetime extending/copy values as appropriate. We rely on
-  /// later optimizations to chew through this traffic. This ensures we can use
-  /// one code base for both OSSA and non-OSSA.
-  OwnershipFixupContext ownershipRAUWHelper;
+  /// External context struct used by \see ownershipRAUWHelper.
+  OwnershipFixupContext ownershipFixupContext;
 
 public:
   SILCombiner(SILOptFunctionBuilder &FuncBuilder, SILBuilder &B,
@@ -134,7 +131,7 @@ class SILCombiner :
               Worklist.add(use->getUser());
             }),
         deBlocks(&B.getFunction()), jPostDomComputer(deBlocks),
-        ownershipRAUWHelper(instModCallbacks, deBlocks, jPostDomComputer) {}
+        ownershipFixupContext(instModCallbacks, deBlocks, jPostDomComputer) {}
 
   bool runOnFunction(SILFunction &F);
 

lib/SILOptimizer/Transforms/CSE.cpp

@@ -583,17 +583,18 @@ class CSE {
 
   DeadEndBlocks &DeadEndBBs;
 
-  OwnershipRAUWHelper &RAUWHelper;
+  OwnershipFixupContext &RAUWFixupContext;
 
   /// The set of calls to lazy property getters which can be replace by a direct
   /// load of the property value.
   llvm::SmallVector<ApplyInst *, 8> lazyPropertyGetters;
 
   CSE(bool RunsOnHighLevelSil, SideEffectAnalysis *SEA,
       SILOptFunctionBuilder &FuncBuilder, DeadEndBlocks &DeadEndBBs,
-      OwnershipRAUWHelper &RAUWHelper)
+      OwnershipFixupContext &RAUWFixupContext)
       : SEA(SEA), FuncBuilder(FuncBuilder), DeadEndBBs(DeadEndBBs),
-        RAUWHelper(RAUWHelper), RunsOnHighLevelSil(RunsOnHighLevelSil) {}
+        RAUWFixupContext(RAUWFixupContext),
+        RunsOnHighLevelSil(RunsOnHighLevelSil) {}
 
   bool processFunction(SILFunction &F, DominanceInfo *DT);
 
@@ -1017,14 +1018,14 @@ bool CSE::processNode(DominanceInfoNode *Node) {
         // extend it here as well
         if (!isa<SingleValueInstruction>(Inst))
           continue;
-        if (!OwnershipRAUWHelper::canFixUpOwnershipForRAUW(
-                cast<SingleValueInstruction>(Inst),
-                cast<SingleValueInstruction>(AvailInst)))
+
+        OwnershipRAUWHelper helper(RAUWFixupContext,
+                                   cast<SingleValueInstruction>(Inst),
+                                   cast<SingleValueInstruction>(AvailInst));
+        if (!helper.isValid())
           continue;
         // Replace SingleValueInstruction using OSSA RAUW here
-        nextI = RAUWHelper.replaceAllUsesAndErase(
-            cast<SingleValueInstruction>(Inst),
-            cast<SingleValueInstruction>(AvailInst));
+        nextI = helper.perform();
         Changed = true;
         ++NumCSE;
         continue;
@@ -1399,8 +1400,7 @@ class SILCSE : public SILFunctionTransform {
     JointPostDominanceSetComputer Computer(DeadEndBBs);
     InstModCallbacks callbacks;
     OwnershipFixupContext FixupCtx{callbacks, DeadEndBBs, Computer};
-    OwnershipRAUWHelper RAUWHelper(FixupCtx);
-    CSE C(RunsOnHighLevelSil, SEA, FuncBuilder, DeadEndBBs, RAUWHelper);
+    CSE C(RunsOnHighLevelSil, SEA, FuncBuilder, DeadEndBBs, FixupCtx);
     bool Changed = false;
 
     // Perform the traditional CSE.