Merge pull request #35302 from atrick/fix-foreachloop

Fix ForEachLoopUnroll use-after-free miscompile.

Author: atrick

lib/SILOptimizer/LoopTransforms/ForEachLoopUnroll.cpp

@@ -88,22 +88,24 @@
 //      alloc_stack %stack
 //      %elem1borrow = begin_borrow %elem1copy
 //      store_borrow %elem1borrow to %stack
-//      end_borrow %elem1borrow
 //      try_apply %body(%stack) normal normalbb1, error errbb1
 //
 //    errbb1(%error : @owned $Error):
-//       br bb2(%error)
+//      end_borrow %elem1borrow
+//      br bb2(%error)
 //
 //    normalbb1(%res : $()):
+//      end_borrow %elem1borrow
 //      %elem2borrow = begin_borrow %elem2copy
 //      store_borrow %elem2borrow to %stack
-//      end_borrow %elem2borrow
 //      try_apply %body(%stack) normal bb1, error errbb2
 //
 //    errbb2(%error : @owned $Error):
-//       br bb2(%error)
+//      end_borrow %elem2borrow
+//      br bb2(%error)
 //
 //    bb1(%res : $()):
+//      end_borrow %elem2borrow
 //      dealloc_stack %stack
 //      ...
 //    bb2(%error : @owned $Error):
@@ -499,12 +501,16 @@ static void unrollForEach(ArrayInfo &arrayInfo, TryApplyInst *forEachCall,
   // A generator for creating a basic block for use as the target of the
   // "error" branch of a try_apply. The error block created here will always
   // jump to the error target of the original forEach.
-  auto errorTargetGenerator = [&](SILBasicBlock *insertionBlock) {
+  auto errorTargetGenerator = [&](SILBasicBlock *insertionBlock,
+                                  SILValue borrowedElem ) {
     SILBasicBlock *newErrorBB = fun->createBasicBlockBefore(insertionBlock);
     SILValue argument = newErrorBB->createPhiArgument(
         errorArgument->getType(), errorArgument->getOwnershipKind());
     // Make the errorBB jump to the error target of the original forEach.
     SILBuilderWithScope builder(newErrorBB, forEachCall);
+    if (borrowedElem) {
+      builder.createEndBorrow(forEachLoc, borrowedElem);
+    }
     builder.createBranch(forEachLoc, errorBB, argument);
     return newErrorBB;
   };
@@ -525,28 +531,34 @@ static void unrollForEach(ArrayInfo &arrayInfo, TryApplyInst *forEachCall,
   //     error blocks jump to the error target of the original forEach call.
   for (uint64_t num = arrayInfo.getNumElements(); num > 0; --num) {
     SILValue elementCopy = elementCopies[num - 1];
+    // Creating the next normal block ends the borrow scope for borrowedElem
+    // from the previous iteration.
     SILBasicBlock *currentBB = num > 1 ? normalTargetGenerator(nextNormalBB)
                                        : forEachCall->getParentBlock();
     SILBuilderWithScope unrollBuilder(currentBB, forEachCall);
+    SILValue borrowedElem;
     if (arrayElementType.isTrivial(*fun)) {
       unrollBuilder.createStore(forEachLoc, elementCopy, allocStack,
                                 StoreOwnershipQualifier::Trivial);
     } else {
       // Borrow the elementCopy and store it in the allocStack. Note that the
       // element's copy is guaranteed to be alive until the array is alive.
       // Therefore it is okay to use a borrow here.
-      SILValue borrowedElem =
-          unrollBuilder.createBeginBorrow(forEachLoc, elementCopy);
+      borrowedElem = unrollBuilder.createBeginBorrow(forEachLoc, elementCopy);
       unrollBuilder.createStoreBorrow(forEachLoc, borrowedElem, allocStack);
-      unrollBuilder.createEndBorrow(forEachLoc, borrowedElem);
+
+      SILBuilderWithScope(&nextNormalBB->front(), forEachCall)
+        .createEndBorrow(forEachLoc, borrowedElem);
     }
-    SILBasicBlock *errorTarget = errorTargetGenerator(nextNormalBB);
+
+    SILBasicBlock *errorTarget =
+      errorTargetGenerator(nextNormalBB, borrowedElem);
     // Note that the substitution map must be empty as we are not supporting
     // elements of address-only type. All elements in the array are guaranteed
     // to be loadable. TODO: generalize this to address-only types.
     unrollBuilder.createTryApply(forEachLoc, forEachBodyClosure,
-                                 SubstitutionMap(), allocStack, nextNormalBB,
-                                 errorTarget);
+                                 SubstitutionMap(), allocStack,
+                                 nextNormalBB, errorTarget);
     nextNormalBB = currentBB;
   }
 

test/SILOptimizer/for_each_loop_unroll_test.sil

@@ -117,23 +117,23 @@ bb2(%39 : @owned $Error):
 // CHECK: [[STACK:%[0-9]+]] = alloc_stack $@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
 // CHECK: [[ELEM1BORROW:%[0-9]+]] = begin_borrow [[ELEM1]]
 // CHECK: store_borrow [[ELEM1BORROW]] to [[STACK]]
-// CHECK: end_borrow [[ELEM1BORROW]]
 // CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error Error, normal [[NORMAL:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
 
 // CHECK: [[ERROR]]([[ERRPARAM:%[0-9]+]] : @owned $Error):
 // CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM]] : $Error)
 
 // CHECK: [[NORMAL]](%{{.*}} : $()):
+// CHECK: end_borrow [[ELEM1BORROW]]
 // CHECK: [[ELEM2BORROW:%[0-9]+]] = begin_borrow [[ELEM2]]
 // CHECK: store_borrow [[ELEM2BORROW]] to [[STACK]]
-// CHECK: end_borrow [[ELEM2BORROW]]
 // CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error Error, normal [[NORMAL2:bb[0-9]+]], error [[ERROR2:bb[0-9]+]]
 
 // CHECK: [[ERROR2]]([[ERRPARAM2:%[0-9]+]] : @owned $Error):
 // CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM2]] : $Error)
 
 // CHECK: [[NORMAL2]](%{{.*}} : $()):
 // CHECK: dealloc_stack [[STACK]]
+// CHECK: end_borrow [[ELEM2BORROW]]
 // 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