Merge pull request #31543 from DougGregor/function-builders-buildarray

[Function builders] Add for...in loop support via buildArray().

Author: DougGregor

include/swift/AST/KnownIdentifiers.def

@@ -31,6 +31,7 @@ IDENTIFIER(Any)
 IDENTIFIER(ArrayLiteralElement)
 IDENTIFIER(atIndexedSubscript)
 IDENTIFIER_(bridgeToObjectiveC)
+IDENTIFIER(buildArray)
 IDENTIFIER(buildBlock)
 IDENTIFIER(buildDo)
 IDENTIFIER(buildEither)

lib/Sema/BuilderTransform.cpp

@@ -741,11 +741,124 @@ class BuilderClosureVisitor
     return visit(caseStmt->getBody());
   }
 
+  VarDecl *visitForEachStmt(ForEachStmt *forEachStmt) {
+    // for...in statements are handled via buildArray(_:); bail out if the
+    // builder does not support it.
+    if (!builderSupports(ctx.Id_buildArray)) {
+      if (!unhandledNode)
+        unhandledNode = forEachStmt;
+      return nullptr;
+    }
+
+    // For-each statements require the Sequence protocol. If we don't have
+    // it (which generally means the standard library isn't loaded), fall
+    // out of the function-builder path entirely to let normal type checking
+    // take care of this.
+    auto sequenceProto = TypeChecker::getProtocol(
+        dc->getASTContext(), forEachStmt->getForLoc(),
+        KnownProtocolKind::Sequence);
+    if (!sequenceProto) {
+      if (!unhandledNode)
+        unhandledNode = forEachStmt;
+      return nullptr;
+    }
+
+    // Generate constraints for the loop header. This also wires up the
+    // types for the patterns.
+    auto target = SolutionApplicationTarget::forForEachStmt(
+        forEachStmt, sequenceProto, dc, /*bindPatternVarsOneWay=*/true);
+    if (cs) {
+      if (cs->generateConstraints(target, FreeTypeVariableBinding::Disallow)) {
+        hadError = true;
+        return nullptr;
+      }
+
+      cs->setSolutionApplicationTarget(forEachStmt, target);
+    }
+
+    // Visit the loop body itself.
+    VarDecl *bodyVar = visit(forEachStmt->getBody());
+    if (!bodyVar)
+      return nullptr;
+
+    // If there's no constraint system, there is nothing left to visit.
+    if (!cs)
+      return nullptr;
+
+    // Form a variable of array type that will capture the result of each
+    // iteration of the loop. We need a fresh type variable to remove the
+    // lvalue-ness of the array variable.
+    SourceLoc loc = forEachStmt->getForLoc();
+    VarDecl *arrayVar = buildVar(loc);
+    Type arrayElementType = cs->createTypeVariable(
+        cs->getConstraintLocator(forEachStmt), 0);
+    cs->addConstraint(
+        ConstraintKind::Equal, cs->getType(bodyVar), arrayElementType,
+        cs->getConstraintLocator(
+          forEachStmt, ConstraintLocator::RValueAdjustment));
+    Type arrayType = ArraySliceType::get(arrayElementType);
+    cs->setType(arrayVar, arrayType);
+
+    // Form an initialization of the array to an empty array literal.
+    Expr *arrayInitExpr = ArrayExpr::create(ctx, loc, { }, { }, loc);
+    cs->setContextualType(
+        arrayInitExpr, TypeLoc::withoutLoc(arrayType), CTP_CannotFail);
+    arrayInitExpr = cs->generateConstraints(arrayInitExpr, dc);
+    if (!arrayInitExpr) {
+      hadError = true;
+      return nullptr;
+    }
+    cs->addConstraint(
+        ConstraintKind::Equal, cs->getType(arrayInitExpr), arrayType,
+        cs->getConstraintLocator(arrayInitExpr));
+
+    // Form a call to Array.append(_:) to add the result of executing each
+    // iteration of the loop body to the array formed above.
+    SourceLoc endLoc = forEachStmt->getEndLoc();
+    auto arrayVarRef = buildVarRef(arrayVar, endLoc);
+    auto arrayAppendRef = new (ctx) UnresolvedDotExpr(
+        arrayVarRef, endLoc, DeclNameRef(ctx.getIdentifier("append")),
+        DeclNameLoc(endLoc), /*implicit=*/true);
+    arrayAppendRef->setFunctionRefKind(FunctionRefKind::SingleApply);
+    auto bodyVarRef = buildVarRef(bodyVar, endLoc);
+    Expr *arrayAppendCall = CallExpr::create(
+        ctx, arrayAppendRef, endLoc, { bodyVarRef } , { Identifier() },
+        { endLoc }, endLoc, /*trailingClosure=*/nullptr, /*implicit=*/true);
+    arrayAppendCall = cs->generateConstraints(arrayAppendCall, dc);
+    if (!arrayAppendCall) {
+      hadError = true;
+      return nullptr;
+    }
+
+    // Form the final call to buildArray(arrayVar) to allow the function
+    // builder to reshape the array into whatever it wants as the result of
+    // the for-each loop.
+    auto finalArrayVarRef = buildVarRef(arrayVar, endLoc);
+    auto buildArrayCall = buildCallIfWanted(
+        endLoc, ctx.Id_buildArray, { finalArrayVarRef }, { Identifier() });
+    assert(buildArrayCall);
+    buildArrayCall = cs->generateConstraints(buildArrayCall, dc);
+    if (!buildArrayCall) {
+      hadError = true;
+      return nullptr;
+    }
+
+    // Form a final variable for the for-each expression itself, which will
+    // be initialized with the call to the function builder's buildArray(_:).
+    auto finalForEachVar = buildVar(loc);
+    cs->setType(finalForEachVar, cs->getType(buildArrayCall));
+    applied.capturedStmts.insert(
+      {forEachStmt, {
+          finalForEachVar,
+          { arrayVarRef, arrayInitExpr, arrayAppendCall, buildArrayCall }}});
+
+    return finalForEachVar;
+  }
+
   CONTROL_FLOW_STMT(Guard)
   CONTROL_FLOW_STMT(While)
   CONTROL_FLOW_STMT(DoCatch)
   CONTROL_FLOW_STMT(RepeatWhile)
-  CONTROL_FLOW_STMT(ForEach)
   CONTROL_FLOW_STMT(Case)
   CONTROL_FLOW_STMT(Break)
   CONTROL_FLOW_STMT(Continue)
@@ -765,6 +878,9 @@ struct FunctionBuilderTarget {
     ReturnValue,
     /// The temporary variable into which the result should be assigned.
     TemporaryVar,
+    /// An expression to evaluate at the end of the block, allowing the update
+    /// of some state from an outer scope.
+    Expression,
   } kind;
 
   /// Captured variable information.
@@ -778,6 +894,10 @@ struct FunctionBuilderTarget {
                                          llvm::TinyPtrVector<Expr *> exprs) {
     return FunctionBuilderTarget{TemporaryVar, {temporaryVar, exprs}};
   }
+
+  static FunctionBuilderTarget forExpression(Expr *expr) {
+    return FunctionBuilderTarget{Expression, { nullptr, { expr }}};
+  }
 };
 
 /// Handles the rewrite of the body of a closure to which a function builder
@@ -882,6 +1002,10 @@ class BuilderClosureRewriter
       assign->setType(TupleType::getEmpty(ctx));
       return assign;
     }
+
+    case FunctionBuilderTarget::Expression:
+      // Execute the expression.
+      return rewriteExpr(capturedExpr);
     }
   }
 
@@ -1181,6 +1305,59 @@ class BuilderClosureRewriter
     return caseStmt;
   }
 
+  Stmt *visitForEachStmt(
+      ForEachStmt *forEachStmt, FunctionBuilderTarget target) {
+    // Translate the for-each loop header.
+    ConstraintSystem &cs = solution.getConstraintSystem();
+    auto forEachTarget =
+        rewriteTarget(*cs.getSolutionApplicationTarget(forEachStmt));
+    if (!forEachTarget)
+      return nullptr;
+
+    const auto &captured = target.captured;
+    auto finalForEachVar = captured.first;
+    auto arrayVarRef = captured.second[0];
+    auto arrayVar = cast<VarDecl>(cast<DeclRefExpr>(arrayVarRef)->getDecl());
+    auto arrayInitExpr = captured.second[1];
+    auto arrayAppendCall = captured.second[2];
+    auto buildArrayCall = captured.second[3];
+
+    // Collect the three steps to initialize the array variable to an
+    // empty array, execute the loop to collect the results of each iteration,
+    // then form the buildArray() call to the write the result.
+    std::vector<ASTNode> outerBodySteps;
+
+    // Step 1: Declare and initialize the array variable.
+    arrayVar->setInterfaceType(solution.simplifyType(cs.getType(arrayVar)));
+    arrayInitExpr = rewriteExpr(arrayInitExpr);
+    declareTemporaryVariable(arrayVar, outerBodySteps, arrayInitExpr);
+
+    // Step 2. Transform the body of the for-each statement. Each iteration
+    // will append the result of executing the loop body to the array.
+    auto body = forEachStmt->getBody();
+    auto capturedBody = takeCapturedStmt(body);
+    auto newBody = cast<BraceStmt>(
+        visitBraceStmt(
+          body,
+          FunctionBuilderTarget::forExpression(arrayAppendCall),
+          FunctionBuilderTarget::forAssign(
+            capturedBody.first, {capturedBody.second.front()})));
+    forEachStmt->setBody(newBody);
+    outerBodySteps.push_back(forEachStmt);
+
+    // Step 3. Perform the buildArray() call to turn the array of results
+    // collected from the iterations into a single value under the control of
+    // the function builder.
+    outerBodySteps.push_back(
+        initializeTarget(
+          FunctionBuilderTarget::forAssign(finalForEachVar, {buildArrayCall})));
+
+    // Form a brace statement to put together the three main steps for the
+    // for-each loop translation outlined above.
+    return BraceStmt::create(
+        ctx, forEachStmt->getStartLoc(), outerBodySteps, newBody->getEndLoc());
+  }
+
 #define UNHANDLED_FUNCTION_BUILDER_STMT(STMT) \
   Stmt *visit##STMT##Stmt(STMT##Stmt *stmt, FunctionBuilderTarget target) { \
     llvm_unreachable("Function builders do not allow statement of kind " \
@@ -1194,7 +1371,6 @@ class BuilderClosureRewriter
   UNHANDLED_FUNCTION_BUILDER_STMT(Defer)
   UNHANDLED_FUNCTION_BUILDER_STMT(DoCatch)
   UNHANDLED_FUNCTION_BUILDER_STMT(RepeatWhile)
-  UNHANDLED_FUNCTION_BUILDER_STMT(ForEach)
   UNHANDLED_FUNCTION_BUILDER_STMT(Break)
   UNHANDLED_FUNCTION_BUILDER_STMT(Continue)
   UNHANDLED_FUNCTION_BUILDER_STMT(Fallthrough)

test/Constraints/function_builder.swift

@@ -41,6 +41,8 @@ struct TupleBuilder {
   static func buildEither<T,U>(second value: U) -> Either<T,U> {
     return .second(value)
   }
+
+  static func buildArray<T>(_ array: [T]) -> [T] { return array }
 }
 
 func tuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) {
@@ -664,3 +666,26 @@ tuplifyWithOpt(true) { c in
   "1"
   3.14159
 }
+
+// Test for-each loops with buildArray.
+// CHECK: testForEach
+// CHECK-SAME: (1, "separator")
+// CHECK-SAME: (2, "separator")
+// CHECK-SAME: (3, "separator")
+// CHECK-SAME: (4, "separator")
+// CHECK-SAME: (5, "separator")
+// CHECK-SAME: (6, "separator")
+// CHECK-SAME: (7, "separator")
+// CHECK-SAME: (8, "separator")
+// CHECK-SAME: (9, "separator")
+// CHECK-SAME: (10, "separator")
+tuplify(true) { c in
+  "testForEach"
+  for i in 0 ..< (c ? 10 : 5) {
+    i + 1
+    "separator"
+  }
+}
+
+
+