OSSA: Rewrite address cloning code to fix issues and fix an infinite loop in the ownership verifier.

include/swift/SIL/MemAccessUtils.h

@@ -1283,7 +1283,9 @@ inline Operand *getAccessProjectionOperand(SingleValueInstruction *svi) {
 /// address must be at operand(0).
 ///
 /// Some of these casts, such as address_to_pointer, may also occur inside of a
-/// formal access. TODO: Add stricter structural guarantee such that these never
+/// formal access.
+///
+/// TODO: Add stricter structural guarantee such that these never
 /// occur within an access. It's important to be able to get the accessed
 /// address without looking though type casts or pointer_to_address [strict],
 /// which we can't do if those operations are behind access projections.
@@ -1497,35 +1499,54 @@ Result AccessUseDefChainVisitor<Impl, Result>::visit(SILValue sourceAddr) {
 
 namespace swift {
 
-/// Clone all projections and casts on the access use-def chain until either the
-/// specified predicate is true or the access base is reached.
+/// Clone all projections and casts on the access use-def chain until the
+/// checkBase predicate returns a valid base.
 ///
-/// This will not clone ref_element_addr or ref_tail_addr because those aren't
-/// part of the access chain.
-template <typename UnaryPredicate>
+/// See comments on the cloneUseDefChain() API.
+template <typename CheckBase>
 class AccessUseDefChainCloner
-    : public AccessUseDefChainVisitor<AccessUseDefChainCloner<UnaryPredicate>,
+    : public AccessUseDefChainVisitor<AccessUseDefChainCloner<CheckBase>,
                                       SILValue> {
-  UnaryPredicate predicate;
+  CheckBase checkBase;
   SILInstruction *insertionPoint;
+  SILValue placeHolder = SILValue();
 
 public:
-  AccessUseDefChainCloner(UnaryPredicate predicate,
-                          SILInstruction *insertionPoint)
-      : predicate(predicate), insertionPoint(insertionPoint) {}
+  AccessUseDefChainCloner(CheckBase checkBase, SILInstruction *insertionPoint)
+      : checkBase(checkBase), insertionPoint(insertionPoint) {}
 
-  // Recursive main entry point
+  // Main entry point
   SILValue cloneUseDefChain(SILValue addr) {
-    if (!predicate(addr))
-      return addr;
+    placeHolder = SILValue();
+    return cloneRecursive(addr);
+  }
+
+  // Secondary entry point to check that cloning will succeed.
+  bool canCloneUseDefChain(SILValue addr) {
+    // Use any valid address as a placeholder. It is innaccessible.
+    placeHolder = addr;
+    return cloneRecursive(addr);
+  }
+
+  SILValue cloneRecursive(SILValue addr) {
+    if (SILValue base = checkBase(addr))
+      return base;
 
     return this->visit(addr);
   }
 
   // Recursively clone an address on the use-def chain.
-  SingleValueInstruction *cloneProjection(SingleValueInstruction *projectedAddr,
-                                          Operand *sourceOper) {
-    SILValue projectedSource = cloneUseDefChain(sourceOper->get());
+  //
+  // Helper for cloneUseDefChain.
+  SILValue cloneProjection(SingleValueInstruction *projectedAddr,
+                           Operand *sourceOper) {
+    SILValue projectedSource = cloneRecursive(sourceOper->get());
+    if (!projectedSource)
+      return nullptr;
+
+    if (placeHolder)
+      return placeHolder;
+
     SILInstruction *clone = projectedAddr->clone(insertionPoint);
     clone->setOperand(sourceOper->getOperandNumber(), projectedSource);
     return cast<SingleValueInstruction>(clone);
@@ -1549,6 +1570,11 @@ class AccessUseDefChainCloner
   }
 
   SILValue visitStorageCast(SingleValueInstruction *cast, Operand *sourceOper) {
+    // The cloner does not currently know how to create compensating
+    // end_borrows or fix mark_dependence operands.
+    if (isa<BeginBorrowInst>(cast) || isa<MarkDependenceInst>(cast))
+      return SILValue();
+
     return cloneProjection(cast, sourceOper);
   }
 
@@ -1558,14 +1584,37 @@ class AccessUseDefChainCloner
   }
 };
 
-template <typename UnaryPredicate>
+/// Clone all projections and casts on the access use-def chain until the
+/// checkBase predicate returns a valid base.
+///
+/// This will not clone ref_element_addr or ref_tail_addr because those aren't
+/// part of the access chain.
+///
+/// CheckBase is a unary predicate that takes the next source address and either
+/// returns a valid SILValue to use as the base of the cloned access path, or an
+/// invalid SILValue to continue cloning.
+///
+/// CheckBase must return a valid SILValue either before attempting to clone the
+/// access base. The most basic valid predicate is:
+///
+///    auto checkBase = [&](SILValue srcAddr) {
+///      return (srcAddr == accessBase) ? srcAddr : SILValue();
+///    }
+template <typename CheckBase>
 SILValue cloneUseDefChain(SILValue addr, SILInstruction *insertionPoint,
-                          UnaryPredicate shouldFollowUse) {
-  return AccessUseDefChainCloner<UnaryPredicate>(shouldFollowUse,
-                                                 insertionPoint)
+                          CheckBase checkBase) {
+  return AccessUseDefChainCloner<CheckBase>(checkBase, insertionPoint)
       .cloneUseDefChain(addr);
 }
 
+/// Analog to cloneUseDefChain to check validity. begin_borrow and
+/// mark_dependence currently cannot be cloned.
+template <typename CheckBase>
+bool canCloneUseDefChain(SILValue addr, CheckBase checkBase) {
+  return AccessUseDefChainCloner<CheckBase>(checkBase, nullptr)
+      .canCloneUseDefChain(addr);
+}
+
 } // end namespace swift
 
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/OwnershipUtils.cpp

@@ -768,7 +768,7 @@ bool InteriorPointerOperand::findTransitiveUsesForAddress(
       }
 
       // If we are the value use, look through it.
-      llvm::copy(mdi->getValue()->getUses(), std::back_inserter(worklist));
+      llvm::copy(mdi->getUses(), std::back_inserter(worklist));
       continue;
     }
 

lib/SILOptimizer/LoopTransforms/LICM.cpp

@@ -1340,11 +1340,21 @@ hoistLoadsAndStores(AccessPath accessPath, SILLoop *loop) {
     }
   }
   assert(storeAddr && "hoistLoadsAndStores requires a store in the loop");
-  SILValue initialAddr = cloneUseDefChain(
-      storeAddr, preheader->getTerminator(), [&](SILValue srcAddr) {
-        // Clone projections until the address dominates preheader.
-        return !DomTree->dominates(srcAddr->getParentBlock(), preheader);
-      });
+  auto checkBase = [&](SILValue srcAddr) {
+    // Clone projections until the address dominates preheader.
+    if (DomTree->dominates(srcAddr->getParentBlock(), preheader))
+      return srcAddr;
+
+    // return an invalid SILValue to continue cloning.
+    return SILValue();
+  };
+  SILValue initialAddr =
+      cloneUseDefChain(storeAddr, preheader->getTerminator(), checkBase);
+  // cloneUseDefChain may currently fail if a begin_borrow or mark_dependence is
+  // in the chain.
+  if (!initialAddr)
+    return;
+
   LoadInst *initialLoad =
       B.createLoad(preheader->getTerminator()->getLoc(), initialAddr,
                    LoadOwnershipQualifier::Unqualified);

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -49,6 +49,12 @@ static llvm::cl::opt<bool> EnableSinkingOwnedForwardingInstToUses(
     llvm::cl::desc("Enable sinking of owened forwarding insts"),
     llvm::cl::init(true), llvm::cl::Hidden);
 
+// Allow disabling general optimization for targetted unit tests.
+static llvm::cl::opt<bool> EnableSILCombineCanonicalize(
+    "sil-combine-canonicalize",
+    llvm::cl::desc("Canonicalization during sil-combine"), llvm::cl::init(true),
+    llvm::cl::Hidden);
+
 //===----------------------------------------------------------------------===//
 //                              Utility Methods
 //===----------------------------------------------------------------------===//
@@ -141,6 +147,9 @@ class SILCombineCanonicalize final : CanonicalizeInstruction {
   }
 
   bool tryCanonicalize(SILInstruction *inst) {
+    if (!EnableSILCombineCanonicalize)
+      return false;
+
     changed = false;
     canonicalize(inst);
     return changed;

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -788,89 +788,6 @@ SILBasicBlock::iterator OwnershipRAUWUtility::perform() {
 //                     Interior Pointer Operand Rebasing
 //===----------------------------------------------------------------------===//
 
-namespace {
-
-/// Clone all projections and casts on the access use-def chain until either the
-/// specified predicate is true or the access base is reached.
-///
-/// This will not clone ref_element_addr or ref_tail_addr because those aren't
-/// part of the access chain.
-class InteriorPointerAddressRebaseUseDefChainCloner
-    : public AccessUseDefChainVisitor<
-          InteriorPointerAddressRebaseUseDefChainCloner, SILValue> {
-  SILValue oldAddressValue;
-  SingleValueInstruction *oldIntPtr;
-  SingleValueInstruction *newIntPtr;
-  SILInstruction *insertPt;
-
-public:
-  InteriorPointerAddressRebaseUseDefChainCloner(
-      SILValue oldAddressValue, SingleValueInstruction *oldIntPtr,
-      SingleValueInstruction *newIntPtr, SILInstruction *insertPt)
-      : oldAddressValue(oldAddressValue), oldIntPtr(oldIntPtr),
-        newIntPtr(newIntPtr), insertPt(insertPt) {}
-
-  // Recursive main entry point
-  SILValue cloneUseDefChain(SILValue currentOldAddr) {
-    // If we have finally hit oldIntPtr, we are done.
-    if (currentOldAddr == oldIntPtr)
-      return currentOldAddr;
-    return this->visit(currentOldAddr);
-  }
-
-  // Recursively clone an address on the use-def chain.
-  SingleValueInstruction *cloneProjection(SingleValueInstruction *projectedAddr,
-                                          Operand *sourceOper) {
-    SILValue sourceOperVal = sourceOper->get();
-    SILValue projectedSource = cloneUseDefChain(sourceOperVal);
-    // If we hit the end of our chain, then make newIntPtr the operand so that
-    // we have successfully rebased.
-    if (sourceOperVal == projectedSource)
-      projectedSource = newIntPtr;
-    SILInstruction *clone = projectedAddr->clone(insertPt);
-    clone->setOperand(sourceOper->getOperandNumber(), projectedSource);
-    return cast<SingleValueInstruction>(clone);
-  }
-
-  SILValue visitBase(SILValue base, AccessedStorage::Kind kind) {
-    assert(false && "access base cannot be cloned");
-    return SILValue();
-  }
-
-  SILValue visitNonAccess(SILValue base) {
-    assert(false && "unknown address root cannot be cloned");
-    return SILValue();
-  }
-
-  SILValue visitPhi(SILPhiArgument *phi) {
-    assert(false && "unexpected phi on access path");
-    return SILValue();
-  }
-
-  SILValue visitStorageCast(SingleValueInstruction *cast, Operand *sourceOper) {
-    // We are ok with cloning storage casts.
-    return cloneProjection(cast, sourceOper);
-  }
-
-  SILValue visitAccessProjection(SingleValueInstruction *projectedAddr,
-                                 Operand *sourceOper) {
-    return cloneProjection(projectedAddr, sourceOper);
-  }
-};
-
-} // namespace
-
-/// \p oldAddressValue is an address rooted in \p oldIntPtr. Clone the use-def
-/// chain from \p oldAddressValue to \p oldIntPtr, but starting from \p
-/// newAddressValue.
-static SILValue cloneInteriorProjectionUseDefChain(
-    SILValue oldAddressValue, SingleValueInstruction *oldIntPtr,
-    SingleValueInstruction *newIntPtr, SILInstruction *insertPt) {
-  InteriorPointerAddressRebaseUseDefChainCloner cloner(
-      oldAddressValue, oldIntPtr, newIntPtr, insertPt);
-  return cloner.cloneUseDefChain(oldAddressValue);
-}
-
 SILBasicBlock::iterator
 OwnershipRAUWHelper::replaceAddressUses(SingleValueInstruction *oldValue,
                                         SILValue newValue) {
@@ -907,19 +824,18 @@ OwnershipRAUWHelper::replaceAddressUses(SingleValueInstruction *oldValue,
   // the access path from newValue to intPtr but upon newIntPtr. Then we make it
   // use newIntPtr.
   auto *intPtrUser = cast<SingleValueInstruction>(intPtr->getUser());
-  SILValue initialAddr = cloneInteriorProjectionUseDefChain(
-      newValue /*address we originally wanted to replace*/,
-      intPtrUser /*the interior pointer of that value*/,
-      newIntPtrUser /*the interior pointer we need to recreate the chain upon*/,
-      oldValue /*insert point*/);
 
-  // If we got back newValue, then we need to set initialAddr to be int ptr
-  // user.
-  if (initialAddr == newValue)
-    initialAddr = newIntPtrUser;
+  // This cloner invocation must match the canCloneUseDefChain check in the
+  // constructor.
+  auto checkBase = [&](SILValue srcAddr) {
+    return (srcAddr == intPtrUser) ? SILValue(newIntPtrUser) : SILValue();
+  };
+  SILValue clonedAddr =
+      cloneUseDefChain(newValue, /*insetPt*/ oldValue, checkBase);
+  assert(clonedAddr != newValue && "expect at least the base to be replaced");
 
   // Now that we have an addr that is setup appropriately, RAUW!
-  return replaceAllUsesAndErase(oldValue, initialAddr, ctx->callbacks);
+  return replaceAllUsesAndErase(oldValue, clonedAddr, ctx->callbacks);
 }
 
 //===----------------------------------------------------------------------===//
@@ -1015,6 +931,17 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
     return;
   }
 
+  // This cloner check must match the later cloner invocation in
+  // replaceAddressUses()
+  auto *intPtrUser = cast<SingleValueInstruction>(intPtr->getUser());
+  auto checkBase = [&](SILValue srcAddr) {
+    return (srcAddr == intPtrUser) ? SILValue(intPtrUser) : SILValue();
+  };
+  if (!canCloneUseDefChain(newValue, checkBase)) {
+    ctx = nullptr;
+    return;
+  }
+
   // For now, just gather up uses
   auto &oldValueUses = ctx->extraAddressFixupInfo.allAddressUsesFromOldValue;
   if (InteriorPointerOperand::findTransitiveUsesForAddress(oldValue,