[inst-simplify] Update for OSSA

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(