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

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) {