Fix SILGen and make it work with all RangeReplaceableCollections

Author: MaxDesiatov

lib/SILGen/SILGenExpr.cpp

@@ -2515,6 +2515,16 @@ RValue RValueEmitter::visitEnumIsCaseExpr(EnumIsCaseExpr *E,
 
 RValue RValueEmitter::visitSingleValueStmtExpr(SingleValueStmtExpr *E,
                                                SGFContext C) {
+  if (E->getStmtKind() == SingleValueStmtExpr::Kind::For) {
+    auto *decl = E->getForExpressionPreamble()->ForAccumulatorDecl;
+    auto *binding = E->getForExpressionPreamble()->ForAccumulatorBinding;
+    SGF.visit(decl);
+    SGF.visit(binding);
+    SGF.emitStmt(E->getStmt());
+
+    return SGF.emitRValueForDecl(E, ConcreteDeclRef(decl), E->getType(), AccessSemantics::Ordinary);
+  }
+
   auto emitStmt = [&]() {
     SGF.emitStmt(E->getStmt());
 

lib/Sema/CSSyntacticElement.cpp

@@ -1503,18 +1503,19 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
     Type elementType = DependentMemberType::get(resultType, sequenceProtocol->getAssociatedType(ctx.Id_Element));
 
     addConstraint(ConstraintKind::Bind, elementTypeVar, elementType, loc);
-    addConstraint(ConstraintKind::Bind, resultType, ArraySliceType::get(elementTypeVar), loc);
+    addConstraint(ConstraintKind::Defaultable, resultType, ArraySliceType::get(elementTypeVar), loc);
 
     auto *binding = E->getForExpressionPreamble()->ForAccumulatorBinding;
 
     auto *initializer = binding->getInit(0);
     auto target = SyntacticElementTarget::forInitialization(initializer, Type(), binding, 0, false);
+    setTargetFor({ binding, 0 }, target);
 
     if (generateConstraints(target)) {
       return true;
     }
 
-    addConstraint(ConstraintKind::Conversion, getType(initializer), resultType, loc);
+    addConstraint(ConstraintKind::Bind, getType(initializer), resultType, loc);
   }
 
   // Propagate the implied result kind from the if/switch expression itself
@@ -2692,6 +2693,18 @@ bool ConstraintSystem::applySolutionToSingleValueStmt(
   if (!stmt || application.hadError)
     return true;
 
+  if (SVE->getStmtKind() == SingleValueStmtExpr::Kind::For) {
+    auto *binding = SVE->getForExpressionPreamble()->ForAccumulatorBinding;
+    auto target = getTargetFor({ binding, 0 }).value();
+
+    auto newTarget = rewriter.rewriteTarget(target);
+    if (!newTarget) {
+      return true;
+    }
+
+    binding->setInit(0, newTarget->getAsExpr());
+  }
+
   SVE->setStmt(stmt);
   return false;
 }

lib/Sema/TypeCheckEffects.cpp

@@ -4036,6 +4036,12 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     ContextScope scope(*this, /*newContext*/ std::nullopt);
     scope.setCoverageForSingleValueStmtExpr();
     SVE->getStmt()->walk(*this);
+
+    if (auto preamble = SVE->getForExpressionPreamble()) {
+      preamble->ForAccumulatorDecl->walk(*this);
+      preamble->ForAccumulatorBinding->walk(*this);
+    }
+
     scope.preserveCoverageFromSingleValueStmtExpr();
     return ShouldNotRecurse;
   }

test/Sema/for-expr.swift

@@ -0,0 +1,12 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature ForExpressions -enable-experimental-feature ThenStatements
+
+func f() -> String {
+  for (i, x) in "hello".enumerated() {
+    then if i % 2 == 0 {
+      x.uppercased()
+    } else  {
+      x
+    }
+  }
+}
+print(f())