Fix logic related to isTriviallyDuplicatable.

lib/SIL/LoopInfo.cpp

@@ -47,6 +47,15 @@ bool SILLoop::canDuplicate(SILInstruction *I) const {
     }
     return true;
   }
+  if (I->isDeallocatingStack()) {
+    SILInstruction *Alloc = nullptr;
+    if (auto *Dealloc = dyn_cast<DeallocStackInst>(I))
+      Alloc = dyn_cast<AllocStackInst>(Dealloc->getOperand());
+    if (auto *Dealloc = dyn_cast<DeallocRefInst>(I))
+      Alloc = dyn_cast<AllocRefInst>(Dealloc->getOperand());
+    // The matching alloc_stack must be in the loop.
+    return Alloc && contains(Alloc);
+  }
 
   // CodeGen can't build ssa for objc methods.
   if (auto *Method = dyn_cast<MethodInst>(I)) {
@@ -71,13 +80,17 @@ bool SILLoop::canDuplicate(SILInstruction *I) const {
     return true;
   }
 
-  if (auto *Dealloc = dyn_cast<DeallocStackInst>(I)) {
-    // The matching alloc_stack must be in the loop.
-    if (auto *Alloc = dyn_cast<AllocStackInst>(Dealloc->getOperand()))
-        return contains(Alloc->getParent());
+  if (isa<ThrowInst>(I))
     return false;
-  }
 
+  // The entire access must be within the loop.
+  if (auto BAI = dyn_cast<BeginAccessInst>(I)) {
+    for (auto *UI : BAI->getUses()) {
+      if (!contains(UI->getUser()))
+        return false;
+    }
+    return true;
+  }
   // The entire coroutine execution must be within the loop.
   // Note that we don't have to worry about the reverse --- a loop which
   // contains an end_apply or abort_apply of an external begin_apply ---
@@ -92,18 +105,11 @@ bool SILLoop::canDuplicate(SILInstruction *I) const {
     return true;
   }
 
-  if (isa<ThrowInst>(I))
-    return false;
-
-  if (isa<BeginAccessInst>(I))
-    return false;
-
   if (isa<DynamicMethodBranchInst>(I))
     return false;
 
-  if (auto *PA = dyn_cast<PartialApplyInst>(I))
-    return !PA->isOnStack();
-
+  // Some special cases above that aren't considered isTriviallyDuplicatable
+  // return true early.
   assert(I->isTriviallyDuplicatable() &&
     "Code here must match isTriviallyDuplicatable in SILInstruction");
   return true;

lib/SIL/SILInstruction.cpp

@@ -1172,9 +1172,9 @@ SILInstruction *SILInstruction::clone(SILInstruction *InsertPt) {
 /// additional handling. It is important to know this information when
 /// you perform such optimizations like e.g. jump-threading.
 bool SILInstruction::isTriviallyDuplicatable() const {
-  if (isa<AllocStackInst>(this) || isa<DeallocStackInst>(this)) {
+  if (isAllocatingStack())
     return false;
-  }
+
   if (auto *ARI = dyn_cast<AllocRefInst>(this)) {
     if (ARI->canAllocOnStack())
       return false;
@@ -1213,9 +1213,6 @@ bool SILInstruction::isTriviallyDuplicatable() const {
   if (isa<DynamicMethodBranchInst>(this))
     return false;
 
-  if (auto *PA = dyn_cast<PartialApplyInst>(this))
-    return !PA->isOnStack();
-
   // If you add more cases here, you should also update SILLoop:canDuplicate.
 
   return true;

lib/SILOptimizer/LoopTransforms/LoopRotate.cpp

@@ -70,6 +70,10 @@ canDuplicateOrMoveToPreheader(SILLoop *loop, SILBasicBlock *preheader,
       invariants.insert(inst);
     } else if (!inst->isTriviallyDuplicatable())
       return false;
+    // It wouldn't make sense to rotate dealloc_stack without also rotating the
+    // alloc_stack, which is covered by isTriviallyDuplicatable.
+    else if (isa<DeallocStackInst>(inst))
+      return false;
     else if (isa<FunctionRefInst>(inst)) {
       moves.push_back(inst);
       invariants.insert(inst);

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -783,7 +783,7 @@ static NullablePtr<EnumElementDecl> getEnumCase(SILValue Val,
 }
 
 static int getThreadingCost(SILInstruction *I) {
-  if (!isa<DeallocStackInst>(I) && !I->isTriviallyDuplicatable())
+  if (!I->isTriviallyDuplicatable())
     return 1000;
 
   // Don't jumpthread function calls.