Skip to content

Commit 945b925

Browse files
authored
Merge pull request #26290 from rjmccall/precheck-function-builders-5.1
[5.1] Function builders: pre-check the original closure body in-place
2 parents b65580c + 7793b9d commit 945b925

File tree

4 files changed

+186
-61
lines changed

4 files changed

+186
-61
lines changed

lib/Sema/BuilderTransform.cpp

Lines changed: 108 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -451,45 +451,6 @@ class BuilderClosureVisitor
451451

452452
} // end anonymous namespace
453453

454-
/// Determine whether the given statement contains a 'return' statement anywhere.
455-
static bool hasReturnStmt(Stmt *stmt) {
456-
class ReturnStmtFinder : public ASTWalker {
457-
public:
458-
bool hasReturnStmt = false;
459-
460-
std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
461-
return { false, expr };
462-
}
463-
464-
std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
465-
// Did we find a 'return' statement?
466-
if (isa<ReturnStmt>(stmt)) {
467-
hasReturnStmt = true;
468-
}
469-
470-
return { !hasReturnStmt, stmt };
471-
}
472-
473-
Stmt *walkToStmtPost(Stmt *stmt) override {
474-
return hasReturnStmt ? nullptr : stmt;
475-
}
476-
477-
std::pair<bool, Pattern*> walkToPatternPre(Pattern *pattern) override {
478-
return { false, pattern };
479-
}
480-
481-
bool walkToDeclPre(Decl *D) override { return false; }
482-
483-
bool walkToTypeLocPre(TypeLoc &TL) override { return false; }
484-
485-
bool walkToTypeReprPre(TypeRepr *T) override { return false; }
486-
};
487-
488-
ReturnStmtFinder finder{};
489-
stmt->walk(finder);
490-
return finder.hasReturnStmt;
491-
}
492-
493454
bool TypeChecker::typeCheckFunctionBuilderFuncBody(FuncDecl *FD,
494455
Type builderType) {
495456
// Try to build a single result expression.
@@ -543,26 +504,36 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
543504
assert(builder && "Bad function builder type");
544505
assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());
545506

546-
// Check the form of this closure to see if we can apply the function-builder
547-
// translation at all.
548-
{
549-
// FIXME: Right now, single-expression closures suppress the function
550-
// builder translation.
551-
if (closure->hasSingleExpressionBody())
552-
return getTypeMatchSuccess();
553-
554-
// The presence of an explicit return suppresses the function builder
555-
// translation.
556-
if (hasReturnStmt(closure->getBody())) {
557-
return getTypeMatchSuccess();
558-
}
507+
// FIXME: Right now, single-expression closures suppress the function
508+
// builder translation.
509+
if (closure->hasSingleExpressionBody())
510+
return getTypeMatchSuccess();
511+
512+
// Pre-check the closure body: pre-check any expressions in it and look
513+
// for return statements.
514+
switch (TC.preCheckFunctionBuilderClosureBody(closure)) {
515+
case FunctionBuilderClosurePreCheck::Okay:
516+
// If the pre-check was okay, apply the function-builder transform.
517+
break;
559518

560-
// Check whether we can apply this function builder.
519+
case FunctionBuilderClosurePreCheck::Error:
520+
// If the pre-check had an error, flag that.
521+
return getTypeMatchFailure(locator);
522+
523+
case FunctionBuilderClosurePreCheck::HasReturnStmt:
524+
// If the closure has a return statement, suppress the transform but
525+
// continue solving the constraint system.
526+
return getTypeMatchSuccess();
527+
}
528+
529+
// Check the form of this closure to see if we can apply the
530+
// function-builder translation at all.
531+
{
532+
// Check whether we can apply this specific function builder.
561533
BuilderClosureVisitor visitor(getASTContext(), this,
562534
/*wantExpr=*/false, builderType);
563535
(void)visitor.visit(closure->getBody());
564536

565-
566537
// If we saw a control-flow statement or declaration that the builder
567538
// cannot handle, we don't have a well-formed function builder application.
568539
if (visitor.unhandledNode) {
@@ -607,8 +578,12 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
607578
/*wantExpr=*/true, builderType);
608579
Expr *singleExpr = visitor.visit(closure->getBody());
609580

610-
if (TC.precheckedClosures.insert(closure).second &&
611-
TC.preCheckExpression(singleExpr, closure))
581+
// We've already pre-checked all the original expressions, but do the
582+
// pre-check to the generated expression just to set up any preconditions
583+
// that CSGen might have.
584+
//
585+
// TODO: just build the AST the way we want it in the first place.
586+
if (TC.preCheckExpression(singleExpr, closure))
612587
return getTypeMatchFailure(locator);
613588

614589
singleExpr = generateConstraints(singleExpr, closure);
@@ -636,3 +611,80 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
636611
locator);
637612
return getTypeMatchSuccess();
638613
}
614+
615+
namespace {
616+
617+
/// Pre-check all the expressions in the closure body.
618+
class PreCheckFunctionBuilderClosure : public ASTWalker {
619+
TypeChecker &TC;
620+
ClosureExpr *Closure;
621+
bool HasReturnStmt = false;
622+
bool HasError = false;
623+
public:
624+
PreCheckFunctionBuilderClosure(TypeChecker &tc, ClosureExpr *closure)
625+
: TC(tc), Closure(closure) {}
626+
627+
FunctionBuilderClosurePreCheck run() {
628+
Stmt *oldBody = Closure->getBody();
629+
630+
Stmt *newBody = oldBody->walk(*this);
631+
632+
// If the walk was aborted, it was because we had a problem of some kind.
633+
assert((newBody == nullptr) == (HasError || HasReturnStmt) &&
634+
"unexpected short-circuit while walking closure body");
635+
if (!newBody) {
636+
if (HasError)
637+
return FunctionBuilderClosurePreCheck::Error;
638+
639+
return FunctionBuilderClosurePreCheck::HasReturnStmt;
640+
}
641+
642+
assert(oldBody == newBody && "pre-check walk wasn't in-place?");
643+
644+
return FunctionBuilderClosurePreCheck::Okay;
645+
}
646+
647+
std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
648+
// Pre-check the expression. If this fails, abort the walk immediately.
649+
// Otherwise, replace the expression with the result of pre-checking.
650+
// In either case, don't recurse into the expression.
651+
if (TC.preCheckExpression(E, /*DC*/ Closure)) {
652+
HasError = true;
653+
return std::make_pair(false, nullptr);
654+
}
655+
656+
return std::make_pair(false, E);
657+
}
658+
659+
std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
660+
// If we see a return statement, abort the walk immediately.
661+
if (isa<ReturnStmt>(S)) {
662+
HasReturnStmt = true;
663+
return std::make_pair(false, nullptr);
664+
}
665+
666+
// Otherwise, recurse into the statement normally.
667+
return std::make_pair(true, S);
668+
}
669+
};
670+
671+
}
672+
673+
FunctionBuilderClosurePreCheck
674+
TypeChecker::preCheckFunctionBuilderClosureBody(ClosureExpr *closure) {
675+
// Single-expression closures should already have been pre-checked.
676+
if (closure->hasSingleExpressionBody())
677+
return FunctionBuilderClosurePreCheck::Okay;
678+
679+
// Check whether we've already done this analysis.
680+
auto it = precheckedFunctionBuilderClosures.find(closure);
681+
if (it != precheckedFunctionBuilderClosures.end())
682+
return it->second;
683+
684+
auto result = PreCheckFunctionBuilderClosure(*this, closure).run();
685+
686+
// Cache the result.
687+
precheckedFunctionBuilderClosures.insert(std::make_pair(closure, result));
688+
689+
return result;
690+
}

lib/Sema/TypeChecker.h

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -538,6 +538,17 @@ enum class CheckedCastContextKind {
538538
EnumElementPattern,
539539
};
540540

541+
enum class FunctionBuilderClosurePreCheck : uint8_t {
542+
/// There were no problems pre-checking the closure.
543+
Okay,
544+
545+
/// There was an error pre-checking the closure.
546+
Error,
547+
548+
/// The closure has a return statement.
549+
HasReturnStmt,
550+
};
551+
541552
/// The Swift type checker, which takes a parsed AST and performs name binding,
542553
/// type checking, and semantic analysis to produce a type-annotated AST.
543554
class TypeChecker final : public LazyResolver {
@@ -669,8 +680,10 @@ class TypeChecker final : public LazyResolver {
669680
/// when executing scripts.
670681
bool InImmediateMode = false;
671682

672-
/// Closure expressions that have already been prechecked.
673-
llvm::SmallPtrSet<ClosureExpr *, 2> precheckedClosures;
683+
/// Closure expressions whose bodies have already been prechecked as
684+
/// part of trying to apply a function builder.
685+
llvm::DenseMap<ClosureExpr *, FunctionBuilderClosurePreCheck>
686+
precheckedFunctionBuilderClosures;
674687

675688
/// A helper to construct and typecheck call to super.init().
676689
///
@@ -1229,6 +1242,16 @@ class TypeChecker final : public LazyResolver {
12291242
/// expression and folding sequence expressions.
12301243
bool preCheckExpression(Expr *&expr, DeclContext *dc);
12311244

1245+
/// Pre-check the body of the given closure, which we are about to
1246+
/// generate constraints for.
1247+
///
1248+
/// This mutates the body of the closure, but only in ways that should be
1249+
/// valid even if we end up not actually applying the function-builder
1250+
/// transform: it just does a normal pre-check of all the expressions in
1251+
/// the closure.
1252+
FunctionBuilderClosurePreCheck
1253+
preCheckFunctionBuilderClosureBody(ClosureExpr *closure);
1254+
12321255
/// \name Name lookup
12331256
///
12341257
/// Routines that perform name lookup.

test/Constraints/function_builder.swift

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,3 +301,52 @@ func testAcceptColorTagged(b: Bool, i: Int, s: String, d: Double) {
301301
}
302302

303303
testAcceptColorTagged(b: true, i: 17, s: "Hello", d: 3.14159)
304+
305+
// rdar://53325810
306+
307+
// Test that we don't have problems with expression pre-checking when
308+
// type-checking an overloaded function-builder call. In particular,
309+
// we need to make sure that expressions in the closure are pre-checked
310+
// before we build constraints for them. Note that top-level expressions
311+
// that need to be rewritten by expression prechecking (such as the operator
312+
// sequences in the boolean conditions and statements below) won't be
313+
// rewritten in the original closure body if we just precheck the
314+
// expressions produced by the function-builder transformation.
315+
struct ForEach1<Data : RandomAccessCollection, Content> {
316+
var data: Data
317+
var content: (Data.Element) -> Content
318+
319+
func show() {
320+
print(content(data.first!))
321+
print(content(data.last!))
322+
}
323+
}
324+
extension ForEach1 where Data.Element: StringProtocol {
325+
// Checking this overload shouldn't trigger inappropriate caching that
326+
// affects checking the next overload.
327+
init(_ data: Data,
328+
@TupleBuilder content: @escaping (Data.Element) -> Content) {
329+
self.init(data: data, content: content)
330+
}
331+
}
332+
extension ForEach1 where Data == Range<Int> {
333+
// This is the overload we actually want.
334+
init(_ data: Data,
335+
@TupleBuilder content: @escaping (Int) -> Content) {
336+
self.init(data: data, content: content)
337+
}
338+
}
339+
let testForEach1 = ForEach1(-10 ..< 10) { i in
340+
"testForEach1"
341+
if i < 0 {
342+
"begin"
343+
i < -5
344+
} else {
345+
i > 5
346+
"end"
347+
}
348+
}
349+
testForEach1.show()
350+
351+
// CHECK: ("testForEach1", main.Either<(Swift.String, Swift.Bool), (Swift.Bool, Swift.String)>.first("begin", true))
352+
// CHECK: ("testForEach1", main.Either<(Swift.String, Swift.Bool), (Swift.Bool, Swift.String)>.second(true, "end"))

test/Constraints/function_builder_diags.swift

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ func testDiags() {
7676
tuplify(true) { _ in
7777
17
7878
for c in name { // expected-error{{closure containing control flow statement cannot be used with function builder 'TupleBuilder'}}
79+
// expected-error@-1 {{use of unresolved identifier 'name'}}
7980
}
8081
}
8182

@@ -97,11 +98,11 @@ func testDiags() {
9798
struct A { }
9899
struct B { }
99100

100-
func overloadedTuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) -> A {
101+
func overloadedTuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) -> A { // expected-note {{found this candidate}}
101102
return A()
102103
}
103104

104-
func overloadedTuplify<T>(_ cond: Bool, @TupleBuilderWithoutIf body: (Bool) -> T) -> B {
105+
func overloadedTuplify<T>(_ cond: Bool, @TupleBuilderWithoutIf body: (Bool) -> T) -> B { // expected-note {{found this candidate}}
105106
return B()
106107
}
107108

@@ -114,7 +115,7 @@ func testOverloading(name: String) {
114115

115116
let _: A = a1
116117

117-
_ = overloadedTuplify(true) { b in
118+
_ = overloadedTuplify(true) { b in // expected-error {{ambiguous use of 'overloadedTuplify(_:body:)'}}
118119
b ? "Hello, \(name)" : "Goodbye"
119120
42
120121
overloadedTuplify(false) {

0 commit comments

Comments
 (0)