for expressions scaffolding, making constraints work

Author: MaxDesiatov

include/swift/AST/Expr.h

@@ -6498,16 +6498,22 @@ class KeyPathDotExpr : public Expr {
   }
 };
 
+struct ForCollectionInit {
+  VarDecl *ForAccumulatorDecl;
+  PatternBindingDecl *ForAccumulatorBinding;
+};
+
 /// An expression that may wrap a statement which produces a single value.
 class SingleValueStmtExpr : public Expr {
 public:
   enum class Kind {
-    If, Switch, Do, DoCatch
+    If, Switch, Do, DoCatch, For
   };
 
 private:
   Stmt *S;
   DeclContext *DC;
+  std::optional<ForCollectionInit> ForExpressionPreamble;
 
   SingleValueStmtExpr(Stmt *S, DeclContext *DC)
       : Expr(ExprKind::SingleValueStmt, /*isImplicit*/ true), S(S), DC(DC) {}
@@ -6572,6 +6578,14 @@ class SingleValueStmtExpr : public Expr {
 
   SourceRange getSourceRange() const;
 
+  std::optional<ForCollectionInit> getForExpressionPreamble() const {
+    return this->ForExpressionPreamble;
+  }
+
+  void setForExpressionPreamble(ForCollectionInit newPreamble) {
+    this->ForExpressionPreamble = newPreamble;
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::SingleValueStmt;
   }

include/swift/AST/KnownIdentifiers.def

@@ -29,6 +29,7 @@ IDENTIFIER(alloc)
 IDENTIFIER(allocWithZone)
 IDENTIFIER(allZeros)
 IDENTIFIER(accumulated)
+IDENTIFIER(append)
 IDENTIFIER(ActorType)
 IDENTIFIER(Any)
 IDENTIFIER(ArrayLiteralElement)

include/swift/Basic/Features.def

@@ -424,6 +424,9 @@ EXPERIMENTAL_FEATURE(ThenStatements, false)
 /// Enable 'do' expressions.
 EXPERIMENTAL_FEATURE(DoExpressions, false)
 
+/// Enable 'for' expressions.
+EXPERIMENTAL_FEATURE(ForExpressions, false)
+
 /// Enable implicitly treating the last expression in a function, closure,
 /// and 'if'/'switch' expression as the result.
 EXPERIMENTAL_FEATURE(ImplicitLastExprResults, false)

include/swift/Sema/SyntacticElementTarget.h

@@ -10,7 +10,8 @@
 //
 //===----------------------------------------------------------------------===//
 //
-//  This file defines the SyntacticElementTarget class.
+//  This file defines the SyntacticElementTarget class (a unit of
+//  type-checking).
 //
 //===----------------------------------------------------------------------===//
 
@@ -59,8 +60,8 @@ struct PackIterationInfo {
 /// within the constraint system.
 using ForEachStmtInfo = TaggedUnion<SequenceIterationInfo, PackIterationInfo>;
 
-/// Describes the target to which a constraint system's solution can be
-/// applied.
+/// Describes the target (a unit of type-checking) to which a constraint
+/// system's solution can be applied.
 class SyntacticElementTarget {
 public:
   enum class Kind {

lib/AST/ASTDumper.cpp

@@ -4496,6 +4496,10 @@ class PrintExpr : public ExprVisitor<PrintExpr, void, Label>,
   void visitSingleValueStmtExpr(SingleValueStmtExpr *E, Label label) {
     printCommon(E, "single_value_stmt_expr", label);
     printDeclContext(E);
+    if (auto preamble = E->getForExpressionPreamble()) {
+      printRec(preamble->ForAccumulatorDecl, Label::optional("for_preamble_accumulator_decl"));
+      printRec(preamble->ForAccumulatorBinding, Label::optional("for_preamble_accumulator_binding"));
+    }
     printRec(E->getStmt(), &E->getDeclContext()->getASTContext(),
              Label::optional("stmt"));
     printFoot();

lib/AST/ASTWalker.cpp

@@ -1385,6 +1385,16 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   Expr *visitKeyPathDotExpr(KeyPathDotExpr *E) { return E; }
 
   Expr *visitSingleValueStmtExpr(SingleValueStmtExpr *E) {
+    if (auto preamble = E->getForExpressionPreamble()) {
+      if (doIt(preamble->ForAccumulatorDecl)) {
+        return nullptr;
+      }
+
+      if (doIt(preamble->ForAccumulatorBinding)) {
+        return nullptr;
+      }
+    }
+
     if (auto *S = doIt(E->getStmt())) {
       E->setStmt(S);
     } else {

lib/AST/Expr.cpp

@@ -2763,6 +2763,8 @@ SingleValueStmtExpr::Kind SingleValueStmtExpr::getStmtKind() const {
     return Kind::Do;
   case StmtKind::DoCatch:
     return Kind::DoCatch;
+  case StmtKind::ForEach:
+    return Kind::For;
   default:
     llvm_unreachable("Unhandled kind!");
   }
@@ -2781,6 +2783,9 @@ SingleValueStmtExpr::getBranches(SmallVectorImpl<Stmt *> &scratch) const {
     return scratch;
   case Kind::DoCatch:
     return cast<DoCatchStmt>(getStmt())->getBranches(scratch);
+  case Kind::For:
+    scratch.push_back(cast<ForEachStmt>(getStmt())->getBody());
+    return scratch;
   }
   llvm_unreachable("Unhandled case in switch!");
 }

lib/AST/FeatureSet.cpp

@@ -116,6 +116,7 @@ UNINTERESTING_FEATURE(RegionBasedIsolation)
 UNINTERESTING_FEATURE(PlaygroundExtendedCallbacks)
 UNINTERESTING_FEATURE(ThenStatements)
 UNINTERESTING_FEATURE(DoExpressions)
+UNINTERESTING_FEATURE(ForExpressions)
 UNINTERESTING_FEATURE(ImplicitLastExprResults)
 UNINTERESTING_FEATURE(RawLayout)
 UNINTERESTING_FEATURE(Embedded)

lib/Sema/CSSyntacticElement.cpp

@@ -1488,8 +1488,34 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
   auto &ctx = getASTContext();
 
   auto *loc = getConstraintLocator(E);
-  Type resultTy = createTypeVariable(loc, /*options*/ 0);
-  setType(E, resultTy);
+  Type resultType = createTypeVariable(loc, /*options*/ 0);
+  setType(E, resultType);
+
+  if (E->getStmtKind() == SingleValueStmtExpr::Kind::For) {
+    auto *rrcProtocol = ctx.getProtocol(KnownProtocolKind::RangeReplaceableCollection);
+    auto *sequenceProtocol = ctx.getProtocol(KnownProtocolKind::Sequence);
+
+    addConstraint(ConstraintKind::ConformsTo,
+                  resultType,
+                  rrcProtocol->getDeclaredInterfaceType(),
+                  loc);
+    Type elementTypeVar = createTypeVariable(loc, /*options*/ 0);
+    Type elementType = DependentMemberType::get(resultType, sequenceProtocol->getAssociatedType(ctx.Id_Element));
+
+    addConstraint(ConstraintKind::Bind, elementTypeVar, elementType, loc);
+    addConstraint(ConstraintKind::Bind, resultType, ArraySliceType::get(elementTypeVar), loc);
+
+    auto *binding = E->getForExpressionPreamble()->ForAccumulatorBinding;
+
+    auto *initializer = binding->getInit(0);
+    auto target = SyntacticElementTarget::forInitialization(initializer, Type(), binding, 0, false);
+
+    if (generateConstraints(target)) {
+      return true;
+    }
+
+    addConstraint(ConstraintKind::Conversion, getType(initializer), resultType, loc);
+  }
 
   // Propagate the implied result kind from the if/switch expression itself
   // into the branches.
@@ -1513,21 +1539,24 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
     auto *loc = getConstraintLocator(
         E, {LocatorPathElt::SingleValueStmtResult(idx), ctpElt});
 
-    ContextualTypeInfo info(resultTy, CTP_SingleValueStmtBranch, loc);
+    ContextualTypeInfo info(resultType, CTP_SingleValueStmtBranch, loc);
     setContextualInfo(result, info);
   }
 
   TypeJoinExpr *join = nullptr;
-  if (branches.empty()) {
-    // If we only have statement branches, the expression is typed as Void. This
-    // should only be the case for 'if' and 'switch' statements that must be
-    // expressions that have branches that all end in a throw, and we'll warn
-    // that we've inferred Void.
-    addConstraint(ConstraintKind::Bind, resultTy, ctx.getVoidType(), loc);
-  } else {
-    // Otherwise, we join the result types for each of the branches.
-    join = TypeJoinExpr::forBranchesOfSingleValueStmtExpr(
-        ctx, resultTy, E, AllocationArena::ConstraintSolver);
+
+  if (E->getStmtKind() != SingleValueStmtExpr::Kind::For) {
+    if (branches.empty()) {
+      // If we only have statement branches, the expression is typed as Void. This
+      // should only be the case for 'if' and 'switch' statements that must be
+      // expressions that have branches that all end in a throw, and we'll warn
+      // that we've inferred Void.
+      addConstraint(ConstraintKind::Bind, resultType, ctx.getVoidType(), loc);
+    } else {
+      // Otherwise, we join the result types for each of the branches.
+      join = TypeJoinExpr::forBranchesOfSingleValueStmtExpr(
+                                                            ctx, resultType, E, AllocationArena::ConstraintSolver);
+    }
   }
 
   // If this is an implied return in a closure, we need to account for the fact
@@ -1568,11 +1597,11 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
       if (auto *closureTy = getClosureTypeIfAvailable(CE)) {
         auto closureResultTy = closureTy->getResult();
         auto *bindToClosure = Constraint::create(
-            *this, ConstraintKind::Bind, resultTy, closureResultTy, loc);
+            *this, ConstraintKind::Bind, resultType, closureResultTy, loc);
         bindToClosure->setFavored();
 
         auto *bindToVoid = Constraint::create(*this, ConstraintKind::Bind,
-                                              resultTy, ctx.getVoidType(), loc);
+                                              resultType, ctx.getVoidType(), loc);
 
         addDisjunctionConstraint({bindToClosure, bindToVoid}, loc);
       }

lib/Sema/PreCheckTarget.cpp

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -11,7 +11,8 @@
 //===----------------------------------------------------------------------===//
 //
 // Pre-checking resolves unqualified name references, type expressions and
-// operators.
+// operators. Target in this context refers to `SyntacticElementTarget`, which
+// is a unit of type-checking.
 //
 //===----------------------------------------------------------------------===//
 
@@ -1190,6 +1191,11 @@ class PreCheckTarget final : public ASTWalker {
   /// For the given statement, mark any valid SingleValueStmtExpr children.
   void markAnyValidSingleValueStmts(Stmt *S);
 
+  /// For the given single value expr that's a `for`-expression, run the
+  /// desugaring transformation. For all other kinds of single value statement
+  /// expressions do nothing.
+  void transformForExpression(SingleValueStmtExpr *E);
+
   PreCheckTarget(DeclContext *dc) : Ctx(dc->getASTContext()), DC(dc) {}
 
 public:
@@ -1386,8 +1392,10 @@ class PreCheckTarget final : public ASTWalker {
     if (auto *assignment = dyn_cast<AssignExpr>(expr))
       markAcceptableDiscardExprs(assignment->getDest());
 
-    if (auto *SVE = dyn_cast<SingleValueStmtExpr>(expr))
+    if (auto *SVE = dyn_cast<SingleValueStmtExpr>(expr)) {
       checkSingleValueStmtExpr(SVE);
+      transformForExpression(SVE);
+    }
 
     return finish(true, expr);
   }
@@ -1700,6 +1708,49 @@ void PreCheckTarget::checkSingleValueStmtExpr(SingleValueStmtExpr *SVE) {
   }
 }
 
+void PreCheckTarget::transformForExpression(SingleValueStmtExpr *SVE) {
+  if (SVE->getStmtKind() != SingleValueStmtExpr::Kind::For) {
+    return;
+  }
+
+  auto *declCtx = SVE->getDeclContext();
+  auto &astCtx = declCtx->getASTContext();
+
+  auto sveLoc = SVE->getLoc();
+
+  auto *varDecl = new(astCtx) VarDecl(false, VarDecl::Introducer::Var, sveLoc,
+                                   astCtx.getIdentifier("$forExpressionResult"), declCtx);
+
+  auto namedPattern = NamedPattern::createImplicit(astCtx, varDecl);
+
+  auto *initFunc = new(astCtx) UnresolvedMemberExpr(sveLoc, DeclNameLoc(), DeclNameRef(DeclBaseName::createConstructor()), true);
+  auto *callExpr = CallExpr::createImplicitEmpty(astCtx, initFunc);
+  auto *initExpr = new(astCtx) UnresolvedMemberChainResultExpr(callExpr, initFunc);
+
+  auto *bindingDecl = PatternBindingDecl::createImplicit(astCtx, StaticSpellingKind::None, namedPattern, initExpr, declCtx);
+
+  SVE->setForExpressionPreamble({ varDecl, bindingDecl });
+
+  // For-expressions always have a single branch.
+  SmallVector<Stmt *, 1> scratch;
+  for (auto *branch : SVE->getBranches(scratch)) {
+    auto *BS = dyn_cast<BraceStmt>(branch);
+    if (!BS)
+      continue;
+
+    auto &result = BS->getElements().back();
+    if (auto stmt = result.dyn_cast<Stmt *>()) {
+      if (auto *then = dyn_cast<ThenStmt>(stmt)) {
+        auto *declRefExpr = new(astCtx) DeclRefExpr(varDecl, DeclNameLoc(), true);
+        auto *dotExpr = new(astCtx) UnresolvedDotExpr(declRefExpr, sveLoc, DeclNameRef(DeclBaseName(astCtx.Id_append)), DeclNameLoc(), true);
+        auto *argumentList = ArgumentList::createImplicit(astCtx, { Argument::unlabeled(then->getResult()) });
+        auto *callExpr = CallExpr::createImplicit(astCtx, dotExpr, argumentList);
+        result = callExpr;
+      }
+    }
+  }
+}
+
 void PreCheckTarget::markAnyValidSingleValueStmts(Expr *E) {
   auto findAssignment = [&]() -> AssignExpr * {
     // Don't consider assignments if we have a parent expression (as otherwise