@@ -451,45 +451,6 @@ class BuilderClosureVisitor
451
451
452
452
} // end anonymous namespace
453
453
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
-
493
454
bool TypeChecker::typeCheckFunctionBuilderFuncBody (FuncDecl *FD,
494
455
Type builderType) {
495
456
// Try to build a single result expression.
@@ -543,26 +504,36 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
543
504
assert (builder && " Bad function builder type" );
544
505
assert (builder->getAttrs ().hasAttribute <FunctionBuilderAttr>());
545
506
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 ;
559
518
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.
561
533
BuilderClosureVisitor visitor (getASTContext (), this ,
562
534
/* wantExpr=*/ false , builderType);
563
535
(void )visitor.visit (closure->getBody ());
564
536
565
-
566
537
// If we saw a control-flow statement or declaration that the builder
567
538
// cannot handle, we don't have a well-formed function builder application.
568
539
if (visitor.unhandledNode ) {
@@ -607,8 +578,12 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
607
578
/* wantExpr=*/ true , builderType);
608
579
Expr *singleExpr = visitor.visit (closure->getBody ());
609
580
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))
612
587
return getTypeMatchFailure (locator);
613
588
614
589
singleExpr = generateConstraints (singleExpr, closure);
@@ -636,3 +611,80 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
636
611
locator);
637
612
return getTypeMatchSuccess ();
638
613
}
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
+ }
0 commit comments