[ome] Make sure to rewire all project_box we insert through the mark_uninitialized project_box to preserve DI behavior.

Once I update DI in a bit, this will no longer be necessary.

rdar://29870610

lib/SILOptimizer/Transforms/OwnershipModelEliminator.cpp

@@ -79,6 +79,7 @@ struct OwnershipModelEliminatorVisitor
   bool visitUnmanagedAutoreleaseValueInst(UnmanagedAutoreleaseValueInst *UAVI);
   bool visitCheckedCastBranchInst(CheckedCastBranchInst *CBI);
   bool visitSwitchEnumInst(SwitchEnumInst *SWI);
+  bool visitProjectBoxInst(ProjectBoxInst *PBI);
 };
 
 } // end anonymous namespace
@@ -246,6 +247,67 @@ bool OwnershipModelEliminatorVisitor::visitSwitchEnumInst(
   return true;
 }
 
+// Since we are threading through copies, we may have situations like:
+//
+// let x = alloc_box $Foo
+// let y = project_box x
+// let z = mark_uninitialized y
+// ... use z ...
+//
+// let y2 = project_box x
+//
+// let x2 = copy_value x
+// let y3 = project_box y
+//
+// We need to move project_box like y2 and y3 to go through z so that DI can
+// reason about them.
+//
+// Once DI is updated for ownership, this can go away.
+bool OwnershipModelEliminatorVisitor::visitProjectBoxInst(ProjectBoxInst *PBI) {
+  // First if our operand is already a mark_uninitialized, then we do not need
+  // to do anything.
+  auto *Use = PBI->getSingleUse();
+  if (Use && isa<MarkUninitializedInst>(Use->getUser())) {
+    return false;
+  }
+
+  // Otherwise, lets try to find the alloc_box.
+  SILValue BoxValue = PBI->getOperand();
+  while (auto *CVI = dyn_cast<CopyValueInst>(BoxValue)) {
+    BoxValue = CVI->getOperand();
+  }
+
+  // We were unable to find the alloc_box. This must be an indirect enum box
+  // pattern.
+  auto *ABI = dyn_cast<AllocBoxInst>(BoxValue);
+  if (!ABI)
+    return false;
+
+  // See if we can find (mark_uninitialized (project_box))
+  SILValue MUI;
+  for (auto *Use : ABI->getUses()) {
+    auto *BoxProjection = dyn_cast<ProjectBoxInst>(Use->getUser());
+    if (!BoxProjection)
+      continue;
+    auto *Op = BoxProjection->getSingleUse();
+    if (!Op || !isa<MarkUninitializedInst>(Op->getUser()))
+      continue;
+    MUI = SILValue(Op->getUser());
+    break;
+  }
+
+  // If we did not find a mark uninitialized inst, then this is not the pattern
+  // that we are looking for.
+  if (!MUI)
+    return false;
+
+  // Ok, we found it. Replace all uses of this project box with the
+  // mark_uninitialized and then erase it.
+  PBI->replaceAllUsesWith(MUI);
+  PBI->eraseFromParent();
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 //                           Top Level Entry Point
 //===----------------------------------------------------------------------===//
@@ -265,7 +327,8 @@ struct OwnershipModelEliminator : SILModuleTransform {
       for (auto &BB : F) {
         for (auto II = BB.begin(), IE = BB.end(); II != IE;) {
           // Since we are going to be potentially removing instructions, we need
-          // to make sure to grab out instruction and increment first.
+          // to make sure to increment our iterator before we perform any
+          // visits.
           SILInstruction *I = &*II;
           ++II;