[ownership] Commit a generic replaceAllUsesAndEraseFixingOwnership api and enable SimplifyInstruction on OSSA.

This is a generic API that when ownership is enabled allows one to replace all
uses of a value with a value with a differing ownership by transforming/lifetime
extending as appropriate.

This API supports all pairings of ownership /except/ replacing a value with
OwnershipKind::None with a value without OwnershipKind::None. This is a more
complex optimization that we do not support today. As a result, we include on
our state struct a helper routine that callers can use to know if the two values
that they want to process can be handled by the algorithm.

My moticiation is to use this to to update InstSimplify and SILCombiner in a
less bug prone way rather than just turn stuff off.

Noting that this transformation inserts ownership instructions, I have made sure
to test this API in two ways:

1. With Mandatory Combiner alone (to make sure it works period).

2. With Mandatory Combiner + Semantic ARC Opts to make sure that we can
   eliminate the extra ownership instructions it inserts.

As one can see from the tests, the optimizer today is able to handle all of
these transforms except one conditional case where I need to eliminate a dead
phi arg. I have a separate branch that hits that today but I have exposed unsafe
behavior in ClosureLifetimeFixup that I need to fix first before I can land
that. I don't want that to stop this PR since I think the current low level ARC
optimizer may be able to help me here since this is a simple transform it does
all of the time.

Author: gottesmm

include/swift/SIL/BasicBlockUtils.h

@@ -89,6 +89,8 @@ class DeadEndBlocks {
     }
     return ReachableBlocks.empty();
   }
+
+  const SILFunction *getFunction() const { return F; }
 };
 
 /// A struct that contains the intermediate state used in computing

include/swift/SIL/OwnershipUtils.h

@@ -170,6 +170,15 @@ struct BorrowingOperand {
     return *this;
   }
 
+  // A set of operators so that a BorrowingOperand can be used like a normal
+  // operand in a light weight way.
+  operator const Operand *() const { return op; }
+  operator Operand *() { return op; }
+  const Operand *operator*() const { return op; }
+  Operand *operator*() { return op; }
+  const Operand *operator->() const { return op; }
+  Operand *operator->() { return op; }
+
   /// If \p op is a borrow introducing operand return it after doing some
   /// checks.
   static Optional<BorrowingOperand> get(Operand *op) {
@@ -425,7 +434,7 @@ struct BorrowedValue {
   ///
   /// NOTE: Scratch space is used internally to this method to store the end
   /// borrow scopes if needed.
-  bool areUsesWithinScope(ArrayRef<Operand *> instructions,
+  bool areUsesWithinScope(ArrayRef<Operand *> uses,
                           SmallVectorImpl<Operand *> &scratchSpace,
                           SmallPtrSetImpl<SILBasicBlock *> &visitedBlocks,
                           DeadEndBlocks &deadEndBlocks) const;
@@ -447,6 +456,24 @@ struct BorrowedValue {
   bool visitInteriorPointerOperands(
       function_ref<void(const InteriorPointerOperand &)> func) const;
 
+  /// Visit all immediate uses of this borrowed value and if any of them are
+  /// reborrows, place them in BorrowingOperand form into \p
+  /// foundReborrows. Returns true if we appended any such reborrows to
+  /// foundReborrows... false otherwise.
+  bool
+  gatherReborrows(SmallVectorImpl<BorrowingOperand> &foundReborrows) const {
+    bool foundAnyReborrows = false;
+    for (auto *op : value->getUses()) {
+      if (auto borrowingOperand = BorrowingOperand::get(op)) {
+        if (borrowingOperand->isReborrow()) {
+          foundReborrows.push_back(*borrowingOperand);
+          foundAnyReborrows = true;
+        }
+      }
+    }
+    return foundAnyReborrows;
+  }
+
 private:
   /// Internal constructor for failable static constructor. Please do not expand
   /// its usage since it assumes the code passed in is well formed.

include/swift/SILOptimizer/Analysis/SimplifyInstruction.h

@@ -30,6 +30,8 @@ class SILInstruction;
 /// analysis of the operands of the instruction, without looking at its uses
 /// (e.g. constant folding).  If a simpler result can be found, it is
 /// returned, otherwise a null SILValue is returned.
+///
+/// This is assumed to implement read-none transformations.
 SILValue simplifyInstruction(SILInstruction *I);
 
 /// Replace an instruction with a simplified result and erase it. If the
@@ -38,9 +40,17 @@ SILValue simplifyInstruction(SILInstruction *I);
 ///
 /// If it is nonnull, eraseNotify will be called before each instruction is
 /// deleted.
+///
+/// If it is nonnull and inst is in OSSA, newInstNotify will be called with each
+/// new instruction inserted to compensate for ownership.
+///
+/// NOTE: When OSSA is enabled this API assumes OSSA is properly formed and will
+/// insert compensating instructions.
 SILBasicBlock::iterator replaceAllSimplifiedUsesAndErase(
     SILInstruction *I, SILValue result,
-    std::function<void(SILInstruction *)> eraseNotify = nullptr);
+    std::function<void(SILInstruction *)> eraseNotify = nullptr,
+    std::function<void(SILInstruction *)> newInstNotify = nullptr,
+    DeadEndBlocks *deadEndBlocks = nullptr);
 
 /// Simplify invocations of builtin operations that may overflow.
 /// All such operations return a tuple (result, overflow_flag).

include/swift/SILOptimizer/Utils/CanonicalizeInstruction.h

@@ -26,6 +26,7 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_CANONICALIZEINSTRUCTION_H
 #define SWIFT_SILOPTIMIZER_UTILS_CANONICALIZEINSTRUCTION_H
 
+#include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILInstruction.h"
 #include "llvm/Support/Debug.h"
@@ -38,8 +39,11 @@ struct CanonicalizeInstruction {
   // May be overriden by passes.
   static constexpr const char *defaultDebugType = "sil-canonicalize";
   const char *debugType = defaultDebugType;
+  DeadEndBlocks &deadEndBlocks;
 
-  CanonicalizeInstruction(const char *passDebugType) {
+  CanonicalizeInstruction(const char *passDebugType,
+                          DeadEndBlocks &deadEndBlocks)
+      : deadEndBlocks(deadEndBlocks) {
 #ifndef NDEBUG
     if (llvm::DebugFlag && !llvm::isCurrentDebugType(debugType))
       debugType = passDebugType;
@@ -48,6 +52,8 @@ struct CanonicalizeInstruction {
 
   virtual ~CanonicalizeInstruction();
 
+  const SILFunction *getFunction() const { return deadEndBlocks.getFunction(); }
+
   /// Rewrite this instruction, based on its operands and uses, into a more
   /// canonical representation.
   ///

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -35,13 +35,13 @@ namespace swift {
 class DominanceInfo;
 template <class T> class NullablePtr;
 
-/// Transform a Use Range (Operand*) into a User Range (SILInstruction*)
+/// Transform a Use Range (Operand*) into a User Range (SILInstruction *)
 using UserTransform = std::function<SILInstruction *(Operand *)>;
 using ValueBaseUserRange =
     TransformRange<iterator_range<ValueBase::use_iterator>, UserTransform>;
 
-inline ValueBaseUserRange
-makeUserRange(iterator_range<ValueBase::use_iterator> range) {
+template <typename Range>
+inline TransformRange<Range, UserTransform> makeUserRange(Range range) {
   auto toUser = [](Operand *operand) { return operand->getUser(); };
   return makeTransformRange(range, UserTransform(toUser));
 }

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -0,0 +1,51 @@
+//===--- OwnershipOptUtils.h ----------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+///
+/// Ownership Utilities that rely on SILOptimizer functionality.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
+#define SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
+
+#include "swift/SIL/OwnershipUtils.h"
+#include "swift/SIL/SILModule.h"
+
+namespace swift {
+
+// Defined in BasicBlockUtils.h
+struct JointPostDominanceSetComputer;
+
+struct OwnershipFixupContext {
+  std::function<void(SILInstruction *)> eraseNotify;
+  std::function<void(SILInstruction *)> newInstNotify;
+  DeadEndBlocks &deBlocks;
+  JointPostDominanceSetComputer &jointPostDomSetComputer;
+
+  SILBasicBlock::iterator
+  replaceAllUsesAndEraseFixingOwnership(SingleValueInstruction *oldValue,
+                                        SILValue newValue);
+
+  /// We can not RAUW all old values with new values.
+  ///
+  /// 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(SingleValueInstruction *oldValue,
+                                       SILValue newValue);
+};
+
+} // namespace swift
+
+#endif

lib/SIL/Utils/BasicBlockUtils.cpp

@@ -387,20 +387,22 @@ void DeadEndBlocks::compute() {
 
 void JointPostDominanceSetComputer::findJointPostDominatingSet(
     SILBasicBlock *dominatingBlock, ArrayRef<SILBasicBlock *> dominatedBlockSet,
-    function_ref<void(SILBasicBlock *)> foundInputBlocksNotInJointPostDomSet,
+    function_ref<void(SILBasicBlock *)> inputBlocksFoundDuringWalk,
     function_ref<void(SILBasicBlock *)> foundJointPostDomSetCompletionBlocks,
-    function_ref<void(SILBasicBlock *)> foundInputBlocksInJointPostDomSet) {
+    function_ref<void(SILBasicBlock *)> inputBlocksInJointPostDomSet) {
   // If our reachable block set is empty, assert. This is most likely programmer
   // error.
   assert(dominatedBlockSet.size() != 0);
 
   // If we have a reachable block set with a single block and that block is
   // dominatingBlock, then we return success since a block post-doms its self so
   // it is already complete.
+  //
+  // NOTE: We do not consider this a visiteed
   if (dominatedBlockSet.size() == 1) {
     if (dominatingBlock == dominatedBlockSet[0]) {
-      if (foundInputBlocksInJointPostDomSet)
-        foundInputBlocksInJointPostDomSet(dominatingBlock);
+      if (inputBlocksInJointPostDomSet)
+        inputBlocksInJointPostDomSet(dominatingBlock);
       return;
     }
   }
@@ -475,16 +477,16 @@ void JointPostDominanceSetComputer::findJointPostDominatingSet(
   // callback.
   sortUnique(reachableInputBlocks);
   for (auto *block : reachableInputBlocks)
-    foundInputBlocksNotInJointPostDomSet(block);
+    inputBlocksFoundDuringWalk(block);
 
   // Then if were asked to find the subset of our input blocks that are in the
   // joint-postdominance set, compute that.
-  if (!foundInputBlocksInJointPostDomSet)
+  if (!inputBlocksInJointPostDomSet)
     return;
 
   // Pass back the reachable input blocks that were not reachable from other
   // input blocks to.
   for (auto *block : dominatedBlockSet)
     if (lower_bound(reachableInputBlocks, block) == reachableInputBlocks.end())
-      foundInputBlocksInJointPostDomSet(block);
+      inputBlocksInJointPostDomSet(block);
 }

lib/SIL/Utils/OwnershipUtils.cpp

@@ -221,9 +221,17 @@ bool BorrowingOperand::visitLocalEndScopeUses(
   case BorrowingOperandKind::TryApply:
   case BorrowingOperandKind::Yield:
     return true;
-  case BorrowingOperandKind::Branch:
+  case BorrowingOperandKind::Branch: {
+    auto *br = cast<BranchInst>(op->getUser());
+    for (auto *use : br->getArgForOperand(op)->getUses())
+      if (use->isLifetimeEnding())
+        if (!func(use))
+          return false;
     return true;
   }
+  }
+
+  llvm_unreachable("Covered switch isn't covered");
 }
 
 void BorrowingOperand::visitBorrowIntroducingUserResults(