Fix creation of end_borrows of store_borrows in ForEachLoopUnroll

end_borrow should not be inserted after dealloc_stack of the destination

Author: meg-gupta

lib/SILOptimizer/LoopTransforms/ForEachLoopUnroll.cpp

@@ -538,6 +538,7 @@ static void unrollForEach(ArrayInfo &arrayInfo, TryApplyInst *forEachCall,
     SILBasicBlock *currentBB = num > 1 ? normalTargetGenerator(nextNormalBB)
                                        : forEachCall->getParentBlock();
     SILBuilderWithScope unrollBuilder(currentBB, forEachCall);
+    SILBuilderWithScope normalBuilder(&nextNormalBB->front(), forEachCall);
     SILValue borrowedElem;
     SILValue addr;
 
@@ -552,9 +553,8 @@ static void unrollForEach(ArrayInfo &arrayInfo, TryApplyInst *forEachCall,
       borrowedElem = unrollBuilder.createBeginBorrow(forEachLoc, elementCopy);
       addr =
           unrollBuilder.createStoreBorrow(forEachLoc, borrowedElem, allocStack);
-      SILBuilderWithScope builder(&nextNormalBB->front(), forEachCall);
-      builder.createEndBorrow(forEachLoc, addr);
-      builder.createEndBorrow(forEachLoc, borrowedElem);
+      normalBuilder.createEndBorrow(forEachLoc, addr);
+      normalBuilder.createEndBorrow(forEachLoc, borrowedElem);
     }
 
     SILBasicBlock *errorTarget =
@@ -566,16 +566,18 @@ static void unrollForEach(ArrayInfo &arrayInfo, TryApplyInst *forEachCall,
     unrollBuilder.createTryApply(forEachLoc, forEachBodyClosure,
                                  SubstitutionMap(), addr, nextNormalBB,
                                  errorTarget);
+
+    if (nextNormalBB == normalBB) {
+      // Dealloc the stack in the normalBB and also in errorBB. Note that every
+      // try_apply created during the unrolling must pass through these blocks.
+      normalBuilder.createDeallocStack(forEachLoc, allocStack);
+    }
     nextNormalBB = currentBB;
   }
-
   // Dealloc the stack in the normalBB and also in errorBB. Note that every
   // try_apply created during the unrolling must pass through these blocks.
-  SILBuilderWithScope(&normalBB->front())
-      .createDeallocStack(forEachLoc, allocStack);
   SILBuilderWithScope(&errorBB->front())
       .createDeallocStack(forEachLoc, allocStack);
-
   // Remove the forEach call as it has now been unrolled.
   removeForEachCall(forEachCall, deleter);
 }

test/SILOptimizer/for_each_loop_unroll_test.sil

@@ -134,8 +134,8 @@ bb2(%39 : @owned $Error):
 // CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM2]] : $any Error)
 
 // CHECK: [[NORMAL2]](%{{.*}} : $()):
-// CHECK: dealloc_stack [[STACK]]
 // CHECK: end_borrow [[ELEM2BORROW]]
+// CHECK: dealloc_stack [[STACK]]
 // Note that the temporary alloc_stack of the array created for the forEach call
 // will be cleaned up when the forEach call is removed.
 // CHECK: destroy_value