Merge pull request #34979 from varungandhi-apple/vg-fix-silgen-for-loop

[SILGen] Branch on result of next() before conversion.

Author: varungandhi-apple

lib/SILGen/SILGenStmt.cpp

@@ -972,20 +972,18 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
   // to hold the results.  This will be initialized on every entry into the loop
   // header and consumed by the loop body. On loop exit, the terminating value
   // will be in the buffer.
-  CanType optTy;
-  if (S->getConvertElementExpr()) {
-    optTy = S->getConvertElementExpr()->getType()->getCanonicalType();
-  } else {
-    optTy = OptionalType::get(S->getSequenceConformance().getTypeWitnessByName(
-                                  S->getSequence()->getType(),
-                                  SGF.getASTContext().Id_Element))
-                ->getCanonicalType();
-  }
+  CanType optTy =
+      OptionalType::get(
+          S->getSequenceConformance().getTypeWitnessByName(
+              S->getSequence()->getType(), SGF.getASTContext().Id_Element))
+          ->getCanonicalType();
   auto &optTL = SGF.getTypeLowering(optTy);
+
   SILValue addrOnlyBuf;
-  ManagedValue nextBufOrValue;
+  bool nextResultTyIsAddressOnly =
+      optTL.isAddressOnly() && SGF.silConv.useLoweredAddresses();
 
-  if (optTL.isAddressOnly() && SGF.silConv.useLoweredAddresses())
+  if (nextResultTyIsAddressOnly)
     addrOnlyBuf = SGF.emitTemporaryAllocation(S, optTL.getLoweredType());
 
   // Create a new basic block and jump into it.
@@ -1028,25 +1026,22 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
                                 SGFContext().withFollowingSideEffects()));
   };
 
+  bool hasElementConversion = S->getElementExpr();
   auto buildElementRValue = [&](SILLocation loc, SGFContext ctx) {
     RValue result;
     result = SGF.emitApplyMethod(
         loc, iteratorNextRef, buildArgumentSource(),
         PreparedArguments(ArrayRef<AnyFunctionType::Param>({})),
-        S->getElementExpr() ? SGFContext() : ctx);
-    if (S->getElementExpr()) {
-      SILGenFunction::OpaqueValueRAII pushOpaqueValue(
-          SGF, S->getElementExpr(),
-          std::move(result).getAsSingleValue(SGF, loc));
-      result = SGF.emitRValue(S->getConvertElementExpr(), ctx);
-    }
+        hasElementConversion ? SGFContext() : ctx);
     return result;
   };
+
+  ManagedValue nextBufOrElement;
   // Then emit the loop destination block.
   //
   // Advance the generator.  Use a scope to ensure that any temporary stack
   // allocations in the subexpression are immediately released.
-  if (optTL.isAddressOnly() && SGF.silConv.useLoweredAddresses()) {
+  if (nextResultTyIsAddressOnly) {
     // Create the initialization outside of the innerForScope so that the
     // innerForScope doesn't clean it up.
     auto nextInit = SGF.useBufferAsTemporary(addrOnlyBuf, optTL);
@@ -1061,16 +1056,23 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
       }
       innerForScope.pop();
     }
-    nextBufOrValue = nextInit->getManagedAddress();
+    nextBufOrElement = nextInit->getManagedAddress();
   } else {
     ArgumentScope innerForScope(SGF, SILLocation(S));
-    nextBufOrValue = innerForScope.popPreservingValue(
+    nextBufOrElement = innerForScope.popPreservingValue(
         buildElementRValue(SILLocation(S), SGFContext())
             .getAsSingleValue(SGF, SILLocation(S)));
   }
 
   SILBasicBlock *failExitingBlock = createBasicBlock();
-  SwitchEnumBuilder switchEnumBuilder(SGF.B, S, nextBufOrValue);
+  SwitchEnumBuilder switchEnumBuilder(SGF.B, S, nextBufOrElement);
+
+  auto convertElementRValue = [&](ManagedValue inputValue, SGFContext ctx) -> ManagedValue {
+    SILGenFunction::OpaqueValueRAII pushOpaqueValue(SGF, S->getElementExpr(),
+                                                    inputValue);
+    return SGF.emitRValue(S->getConvertElementExpr(), ctx)
+        .getAsSingleValue(SGF, SILLocation(S));
+  };
 
   switchEnumBuilder.addOptionalSomeCase(
       createBasicBlock(), loopDest.getBlock(),
@@ -1096,13 +1098,22 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
           //
           // *NOTE* If we do not have an address only value, then inputValue is
           // *already properly unwrapped.
-          if (optTL.isAddressOnly() && SGF.silConv.useLoweredAddresses()) {
+          SGFContext loopVarCtx{initLoopVars.get()};
+          if (nextResultTyIsAddressOnly) {
             inputValue = SGF.emitUncheckedGetOptionalValueFrom(
-                S, inputValue, optTL, SGFContext(initLoopVars.get()));
+                S, inputValue, optTL,
+                hasElementConversion ? SGFContext() : loopVarCtx);
           }
 
+          CanType optConvertedTy = optTy;
+          if (hasElementConversion) {
+            inputValue = convertElementRValue(inputValue, loopVarCtx);
+            optConvertedTy =
+                OptionalType::get(S->getConvertElementExpr()->getType())
+                    ->getCanonicalType();
+          }
           if (!inputValue.isInContext())
-            RValue(SGF, S, optTy.getOptionalObjectType(), inputValue)
+            RValue(SGF, S, optConvertedTy.getOptionalObjectType(), inputValue)
                 .forwardInto(SGF, S, initLoopVars.get());
 
           // Now that the pattern has been initialized, check any where

lib/Sema/CSApply.cpp

@@ -8039,13 +8039,13 @@ static Optional<SolutionApplicationTarget> applySolutionToForEachStmt(
   // Convert that Optional<Element> value to the type of the pattern.
   auto optPatternType = OptionalType::get(forEachStmtInfo.initType);
   if (!optPatternType->isEqual(nextResultType)) {
-    OpaqueValueExpr *elementExpr =
-        new (ctx) OpaqueValueExpr(stmt->getInLoc(), nextResultType,
-                                  /*isPlaceholder=*/true);
+    OpaqueValueExpr *elementExpr = new (ctx) OpaqueValueExpr(
+        stmt->getInLoc(), nextResultType->getOptionalObjectType(),
+        /*isPlaceholder=*/true);
     Expr *convertElementExpr = elementExpr;
     if (TypeChecker::typeCheckExpression(
             convertElementExpr, dc,
-            /*contextualInfo=*/{optPatternType, CTP_CoerceOperand})
+            /*contextualInfo=*/{forEachStmtInfo.initType, CTP_CoerceOperand})
             .isNull()) {
       return None;
     }

test/Profiler/pgo_foreach.swift

@@ -43,7 +43,7 @@ public func guessForEach1(x: Int32) -> Int32 {
 // IR-OPT-LABEL: define swiftcc i32 @$s9pgo_foreach10guessWhiles5Int32VAD1x_tF
 
 public func guessForEach2(x: Int32) -> Int32 {
-  // SIL: switch_enum {{.*}} : $Optional<(String, Int32)>, case #Optional.some!enumelt: {{.*}} !case_count(168), case #Optional.none!enumelt: {{.*}} !case_count(42)
+  // SIL: switch_enum {{.*}} : $Optional<(key: String, value: Int32)>, case #Optional.some!enumelt: {{.*}} !case_count(168), case #Optional.none!enumelt: {{.*}} !case_count(42)
 
   var ret : Int32 = 0
   let names = ["John" : Int32(1), "Paul" : Int32(2), "George" : Int32(3), "Ringo" : Int32(x)]

test/SILGen/foreach.swift

@@ -603,3 +603,32 @@ func genericFuncWithConversion<T: C>(list : [T]) {
     print(item)
   }
 }
+
+// SR-8688: Check that branch on result of next() precedes optional injection.
+// If we branch on the converted result of next(), the loop won't terminate.
+//
+// CHECK-LABEL: sil hidden [ossa] @$s7foreach32injectForEachElementIntoOptionalyySaySiGF
+// CHECK: [[NEXT_RESULT:%.*]] = load [trivial] {{.*}} : $*Optional<Int>
+// CHECK: switch_enum [[NEXT_RESULT]] : $Optional<Int>, case #Optional.some!enumelt: [[BB_SOME:bb.*]], case
+// CHECK: [[BB_SOME]]([[X_PRE_BINDING:%.*]] : $Int):
+// CHECK: [[X_BINDING:%.*]] = enum $Optional<Int>, #Optional.some!enumelt, [[X_PRE_BINDING]] : $Int
+// CHECK: debug_value [[X_BINDING]] : $Optional<Int>, let, name "x"
+func injectForEachElementIntoOptional(_ xs: [Int]) {
+  for x : Int? in xs {}
+}
+
+// SR-8688: Check that branch on result of next() precedes optional injection.
+// If we branch on the converted result of next(), the loop won't terminate.
+//
+// CHECK-LABEL: sil hidden [ossa] @$s7foreach32injectForEachElementIntoOptionalyySayxGlF
+// CHECK: copy_addr [take] [[NEXT_RESULT:%.*]] to [initialization] [[NEXT_RESULT_COPY:%.*]] : $*Optional<T>
+// CHECK: switch_enum_addr [[NEXT_RESULT_COPY]] : $*Optional<T>, case #Optional.some!enumelt: [[BB_SOME:bb.*]], case
+// CHECK: [[BB_SOME]]:
+// CHECK: [[X_BINDING:%.*]] = alloc_stack $Optional<T>, let, name "x"
+// CHECK: [[ADDR:%.*]] = unchecked_take_enum_data_addr [[NEXT_RESULT_COPY]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK: [[X_ADDR:%.*]] = init_enum_data_addr [[X_BINDING]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK: copy_addr [take] [[ADDR]] to [initialization] [[X_ADDR]] : $*T
+// CHECK: inject_enum_addr [[X_BINDING]] : $*Optional<T>, #Optional.some!enumelt
+func injectForEachElementIntoOptional<T>(_ xs: [T]) {
+  for x : T? in xs {}
+}