Function builders: pre-check the original closure body in-place.

Prior to this patch, we pre-checked the result of applying the function-builder transformation, but only when we hadn't already pre-checked the closure before.  This causes two problems to arise when the transformation is applied to the same closure along multiple branches of a disjunction.  The first is that any expressions that are synthesized by the transformation will not be pre-checked the second time through, which is a problem if we try to apply different builder types to the same closure (we do cache expressions for identical builder types).  The second is that the pre-check will rewrite sub-expressions in place *in the synthesized expression*, which means that top-level expressions in the original closure body (including `if` conditions) that are now nested in the synthesized expression will not be rewritten in the original closure and therefore will be encountered in their raw state the second time through.

This patch causes all expressions in the original closure body to be pre-checked before doing any other work.  We then pre-check the synthesized expression immediately before generating constraints for it in order to set up the AST appropriately for CSGen; this could be skipped if we just synthesized expressions the way that CSGen wants them, but that seems to be somewhat involved.

Pre-checking is safe to apply to an expression multiple times, so it's
fine if we take this path and then decide not to use a function builder.

I've also merged the check for `return` statements into this same walk, which was convenient.

Fixes rdar://53325810 at least, and probably also some bugs with applying different function builders to the same closure.

Author: rjmccall

lib/Sema/BuilderTransform.cpp

@@ -451,45 +451,6 @@ class BuilderClosureVisitor
 
 } // end anonymous namespace
 
-/// Determine whether the given statement contains a 'return' statement anywhere.
-static bool hasReturnStmt(Stmt *stmt) {
-  class ReturnStmtFinder : public ASTWalker {
-  public:
-    bool hasReturnStmt = false;
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      return { false, expr };
-    }
-
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
-      // Did we find a 'return' statement?
-      if (isa<ReturnStmt>(stmt)) {
-        hasReturnStmt = true;
-      }
-
-      return { !hasReturnStmt, stmt };
-    }
-
-    Stmt *walkToStmtPost(Stmt *stmt) override {
-      return hasReturnStmt ? nullptr : stmt;
-    }
-
-    std::pair<bool, Pattern*> walkToPatternPre(Pattern *pattern) override {
-      return { false, pattern };
-    }
-
-    bool walkToDeclPre(Decl *D) override { return false; }
-
-    bool walkToTypeLocPre(TypeLoc &TL) override { return false; }
-
-    bool walkToTypeReprPre(TypeRepr *T) override { return false; }
-  };
-
-  ReturnStmtFinder finder{};
-  stmt->walk(finder);
-  return finder.hasReturnStmt;
-}
-
 BraceStmt *
 TypeChecker::applyFunctionBuilderBodyTransform(FuncDecl *FD,
                                                BraceStmt *body,
@@ -520,26 +481,36 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
   assert(builder && "Bad function builder type");
   assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());
 
-  // Check the form of this closure to see if we can apply the function-builder
-  // translation at all.
-  {
-    // FIXME: Right now, single-expression closures suppress the function
-    // builder translation.
-    if (closure->hasSingleExpressionBody())
-      return getTypeMatchSuccess();
-
-    // The presence of an explicit return suppresses the function builder
-    // translation.
-    if (hasReturnStmt(closure->getBody())) {
-      return getTypeMatchSuccess();
-    }
+  // FIXME: Right now, single-expression closures suppress the function
+  // builder translation.
+  if (closure->hasSingleExpressionBody())
+    return getTypeMatchSuccess();
+
+  // Pre-check the closure body: pre-check any expressions in it and look
+  // for return statements.
+  switch (TC.preCheckFunctionBuilderClosureBody(closure)) {
+  case FunctionBuilderClosurePreCheck::Okay:
+    // If the pre-check was okay, apply the function-builder transform.
+    break;
 
-    // Check whether we can apply this function builder.
+  case FunctionBuilderClosurePreCheck::Error:
+    // If the pre-check had an error, flag that.
+    return getTypeMatchFailure(locator);
+
+  case FunctionBuilderClosurePreCheck::HasReturnStmt:
+    // If the closure has a return statement, suppress the transform but
+    // continue solving the constraint system.
+    return getTypeMatchSuccess();
+  }
+
+  // Check the form of this closure to see if we can apply the
+  // function-builder translation at all.
+  {
+    // Check whether we can apply this specific function builder.
     BuilderClosureVisitor visitor(getASTContext(), this,
                                   /*wantExpr=*/false, builderType);
     (void)visitor.visit(closure->getBody());
 
-
     // If we saw a control-flow statement or declaration that the builder
     // cannot handle, we don't have a well-formed function builder application.
     if (visitor.unhandledNode) {
@@ -584,8 +555,12 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
                                 /*wantExpr=*/true, builderType);
   Expr *singleExpr = visitor.visit(closure->getBody());
 
-  if (TC.precheckedClosures.insert(closure).second &&
-      TC.preCheckExpression(singleExpr, closure))
+  // We've already pre-checked all the original expressions, but do the
+  // pre-check to the generated expression just to set up any preconditions
+  // that CSGen might have.
+  //
+  // TODO: just build the AST the way we want it in the first place.
+  if (TC.preCheckExpression(singleExpr, closure))
     return getTypeMatchFailure(locator);
 
   singleExpr = generateConstraints(singleExpr, closure);
@@ -613,3 +588,80 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
                 locator);
   return getTypeMatchSuccess();
 }
+
+namespace {
+
+/// Pre-check all the expressions in the closure body.
+class PreCheckFunctionBuilderClosure : public ASTWalker {
+  TypeChecker &TC;
+  ClosureExpr *Closure;
+  bool HasReturnStmt = false;
+  bool HasError = false;
+public:
+  PreCheckFunctionBuilderClosure(TypeChecker &tc, ClosureExpr *closure)
+    : TC(tc), Closure(closure) {}
+
+  FunctionBuilderClosurePreCheck run() {
+    Stmt *oldBody = Closure->getBody();
+
+    Stmt *newBody = oldBody->walk(*this);
+
+    // If the walk was aborted, it was because we had a problem of some kind.
+    assert((newBody == nullptr) == (HasError || HasReturnStmt) &&
+           "unexpected short-circuit while walking closure body");
+    if (!newBody) {
+      if (HasError)
+        return FunctionBuilderClosurePreCheck::Error;
+
+      return FunctionBuilderClosurePreCheck::HasReturnStmt;
+    }
+
+    assert(oldBody == newBody && "pre-check walk wasn't in-place?");
+
+    return FunctionBuilderClosurePreCheck::Okay;
+  }
+
+  std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+    // Pre-check the expression.  If this fails, abort the walk immediately.
+    // Otherwise, replace the expression with the result of pre-checking.
+    // In either case, don't recurse into the expression.
+    if (TC.preCheckExpression(E, /*DC*/ Closure)) {
+      HasError = true;
+      return std::make_pair(false, nullptr);
+    }
+
+    return std::make_pair(false, E);
+  }
+
+  std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+    // If we see a return statement, abort the walk immediately.
+    if (isa<ReturnStmt>(S)) {
+      HasReturnStmt = true;
+      return std::make_pair(false, nullptr);
+    }
+
+    // Otherwise, recurse into the statement normally.
+    return std::make_pair(true, S);
+  }
+};
+
+}
+
+FunctionBuilderClosurePreCheck
+TypeChecker::preCheckFunctionBuilderClosureBody(ClosureExpr *closure) {
+  // Single-expression closures should already have been pre-checked.
+  if (closure->hasSingleExpressionBody())
+    return FunctionBuilderClosurePreCheck::Okay;
+
+  // Check whether we've already done this analysis.
+  auto it = precheckedFunctionBuilderClosures.find(closure);
+  if (it != precheckedFunctionBuilderClosures.end())
+    return it->second;
+
+  auto result = PreCheckFunctionBuilderClosure(*this, closure).run();
+
+  // Cache the result.
+  precheckedFunctionBuilderClosures.insert(std::make_pair(closure, result));
+
+  return result;
+}

lib/Sema/TypeChecker.h

@@ -529,6 +529,17 @@ enum class CheckedCastContextKind {
   EnumElementPattern,
 };
 
+enum class FunctionBuilderClosurePreCheck : uint8_t {
+  /// There were no problems pre-checking the closure.
+  Okay,
+
+  /// There was an error pre-checking the closure.
+  Error,
+
+  /// The closure has a return statement.
+  HasReturnStmt,
+};
+
 /// The Swift type checker, which takes a parsed AST and performs name binding,
 /// type checking, and semantic analysis to produce a type-annotated AST.
 class TypeChecker final : public LazyResolver {
@@ -617,8 +628,10 @@ class TypeChecker final : public LazyResolver {
   /// when executing scripts.
   bool InImmediateMode = false;
 
-  /// Closure expressions that have already been prechecked.
-  llvm::SmallPtrSet<ClosureExpr *, 2> precheckedClosures;
+  /// Closure expressions whose bodies have already been prechecked as
+  /// part of trying to apply a function builder.
+  llvm::DenseMap<ClosureExpr *, FunctionBuilderClosurePreCheck>
+    precheckedFunctionBuilderClosures;
 
   TypeChecker(ASTContext &Ctx);
   friend class ASTContext;
@@ -1169,6 +1182,16 @@ class TypeChecker final : public LazyResolver {
   /// expression and folding sequence expressions.
   bool preCheckExpression(Expr *&expr, DeclContext *dc);
 
+  /// Pre-check the body of the given closure, which we are about to
+  /// generate constraints for.
+  ///
+  /// This mutates the body of the closure, but only in ways that should be
+  /// valid even if we end up not actually applying the function-builder
+  /// transform: it just does a normal pre-check of all the expressions in
+  /// the closure.
+  FunctionBuilderClosurePreCheck
+  preCheckFunctionBuilderClosureBody(ClosureExpr *closure);
+
   /// \name Name lookup
   ///
   /// Routines that perform name lookup.

test/Constraints/function_builder.swift

@@ -301,3 +301,52 @@ func testAcceptColorTagged(b: Bool, i: Int, s: String, d: Double) {
 }
 
 testAcceptColorTagged(b: true, i: 17, s: "Hello", d: 3.14159)
+
+// rdar://53325810
+
+// Test that we don't have problems with expression pre-checking when
+// type-checking an overloaded function-builder call.  In particular,
+// we need to make sure that expressions in the closure are pre-checked
+// before we build constraints for them.  Note that top-level expressions
+// that need to be rewritten by expression prechecking (such as the operator
+// sequences in the boolean conditions and statements below) won't be
+// rewritten in the original closure body if we just precheck the
+// expressions produced by the function-builder transformation.
+struct ForEach1<Data : RandomAccessCollection, Content> {
+  var data: Data
+  var content: (Data.Element) -> Content
+
+  func show() {
+    print(content(data.first!))
+    print(content(data.last!))
+  }
+}
+extension ForEach1 where Data.Element: StringProtocol {
+  // Checking this overload shouldn't trigger inappropriate caching that
+  // affects checking the next overload.
+  init(_ data: Data,
+       @TupleBuilder content: @escaping (Data.Element) -> Content) {
+    self.init(data: data, content: content)
+  }
+}
+extension ForEach1 where Data == Range<Int> {
+  // This is the overload we actually want.
+  init(_ data: Data,
+       @TupleBuilder content: @escaping (Int) -> Content) {
+    self.init(data: data, content: content)
+  }
+}
+let testForEach1 = ForEach1(-10 ..< 10) { i in
+  "testForEach1"
+  if i < 0 {
+    "begin"
+    i < -5
+  } else {
+    i > 5
+    "end"
+  }
+}
+testForEach1.show()
+
+// CHECK: ("testForEach1", main.Either<(Swift.String, Swift.Bool), (Swift.Bool, Swift.String)>.first("begin", true))
+// CHECK: ("testForEach1", main.Either<(Swift.String, Swift.Bool), (Swift.Bool, Swift.String)>.second(true, "end"))

test/Constraints/function_builder_diags.swift

@@ -76,6 +76,7 @@ func testDiags() {
   tuplify(true) { _ in
     17
     for c in name { // expected-error{{closure containing control flow statement cannot be used with function builder 'TupleBuilder'}}
+    // expected-error@-1 {{use of unresolved identifier 'name'}}
     }
   }
 
@@ -97,11 +98,11 @@ func testDiags() {
 struct A { }
 struct B { }
 
-func overloadedTuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) -> A {
+func overloadedTuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) -> A { // expected-note {{found this candidate}}
   return A()
 }
 
-func overloadedTuplify<T>(_ cond: Bool, @TupleBuilderWithoutIf body: (Bool) -> T) -> B {
+func overloadedTuplify<T>(_ cond: Bool, @TupleBuilderWithoutIf body: (Bool) -> T) -> B { // expected-note {{found this candidate}}
   return B()
 }
 
@@ -114,7 +115,7 @@ func testOverloading(name: String) {
 
   let _: A = a1
 
-  _ = overloadedTuplify(true) { b in
+  _ = overloadedTuplify(true) { b in // expected-error {{ambiguous use of 'overloadedTuplify(_:body:)'}}
     b ? "Hello, \(name)" : "Goodbye"
     42
     overloadedTuplify(false) {