[MiscDiagnostics] Walk into the body of a non single-statement closure

if the closure had a function builder transform applied.

This way, function builder closures can have syntactic restrictions
diagnosed the same way as other expressions.

Author: hborla

include/swift/AST/ASTWalker.h

@@ -20,6 +20,7 @@ namespace swift {
 
 class Decl;
 class Expr;
+class ClosureExpr;
 class ModuleDecl;
 class Stmt;
 class Pattern;
@@ -219,7 +220,13 @@ class ASTWalker {
   /// For work that is performed for every top-level expression, this should
   /// be overridden to return false, to avoid duplicating work or visiting
   /// bodies of closures that have not yet been type checked.
-  virtual bool shouldWalkIntoNonSingleExpressionClosure() { return true; }
+  virtual bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *) {
+    return true;
+  }
+
+  /// This method configures whether the walker should visit the body of a
+  /// TapExpr.
+  virtual bool shouldWalkIntoTapExpression() { return true; }
 
   /// This method configures whether the walker should exhibit the legacy
   /// behavior where accessors appear as peers of their storage, rather

lib/AST/ASTWalker.cpp

@@ -774,7 +774,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       return nullptr;
     }
 
-    if (!Walker.shouldWalkIntoNonSingleExpressionClosure())
+    if (!Walker.shouldWalkIntoNonSingleExpressionClosure(expr))
       return expr;
 
     // Handle other closures.
@@ -1104,7 +1104,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       }
     }
 
-    if (!Walker.shouldWalkIntoNonSingleExpressionClosure())
+    if (!Walker.shouldWalkIntoTapExpression())
       return E;
 
     if (auto oldBody = E->getBody()) {

lib/Sema/BuilderTransform.cpp

@@ -627,7 +627,7 @@ class BuilderClosureRewriter
   const Solution &solution;
   DeclContext *dc;
   AppliedBuilderTransform builderTransform;
-  std::function<Expr *(Expr *)> rewriteExprFn;
+  std::function<Expr *(Expr *)> rewriteExpr;
   std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType;
 
   /// Retrieve the temporary variable that will be used to capture the
@@ -729,13 +729,6 @@ class BuilderClosureRewriter
     elements.push_back(pbd);
   }
 
-  Expr *rewriteExpr(Expr *expr) {
-    Expr *result = rewriteExprFn(expr);
-    if (result)
-      performSyntacticExprDiagnostics(expr, dc, /*isExprStmt=*/false);
-    return result;
-  }
-
 public:
   BuilderClosureRewriter(
       const Solution &solution,
@@ -745,7 +738,7 @@ class BuilderClosureRewriter
       std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType
     ) : ctx(solution.getConstraintSystem().getASTContext()),
         solution(solution), dc(dc), builderTransform(builderTransform),
-        rewriteExprFn(rewriteExpr),
+        rewriteExpr(rewriteExpr),
         coerceToType(coerceToType){ }
 
   Stmt *visitBraceStmt(BraceStmt *braceStmt, FunctionBuilderTarget target,

lib/Sema/MiscDiagnostics.cpp

@@ -83,17 +83,23 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
     DiagnoseWalker(const DeclContext *DC, bool isExprStmt)
       : IsExprStmt(isExprStmt), Ctx(DC->getASTContext()), DC(DC) {}
 
-    // Not interested in going outside a basic expression.
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
-      return { false, S };
-    }
     std::pair<bool, Pattern*> walkToPatternPre(Pattern *P) override {
       return { false, P };
     }
-    bool walkToDeclPre(Decl *D) override { return false; }
+
+    bool walkToDeclPre(Decl *D) override {
+      if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
+        return closure->hasAppliedFunctionBuilder();
+      return false;
+    }
+
     bool walkToTypeReprPre(TypeRepr *T) override { return true; }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       // See through implicit conversions of the expression.  We want to be able
@@ -1314,7 +1320,11 @@ static void diagRecursivePropertyAccess(const Expr *E, const DeclContext *DC) {
              cast<VarDecl>(DRE->getDecl())->isSelfParameter();
     }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       Expr *subExpr;
@@ -1478,10 +1488,16 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
 
     // Don't walk into nested decls.
     bool walkToDeclPre(Decl *D) override {
+      if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
+        return closure->hasAppliedFunctionBuilder();
       return false;
     }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       if (auto *CE = dyn_cast<AbstractClosureExpr>(E)) {
@@ -3302,7 +3318,11 @@ static void checkStmtConditionTrailingClosure(ASTContext &ctx, const Expr *E) {
   public:
     DiagnoseWalker(ASTContext &ctx) : Ctx(ctx) { }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       switch (E->getKind()) {
@@ -3447,7 +3467,11 @@ class ObjCSelectorWalker : public ASTWalker {
   ObjCSelectorWalker(const DeclContext *dc, Type selectorTy)
     : Ctx(dc->getASTContext()), DC(dc), SelectorTy(selectorTy) { }
 
-  bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+  bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+    return expr->hasAppliedFunctionBuilder();
+  }
+
+  bool shouldWalkIntoTapExpression() override { return false; }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
     auto *stringLiteral = dyn_cast<StringLiteralExpr>(expr);
@@ -4164,7 +4188,11 @@ static void diagnoseUnintendedOptionalBehavior(const Expr *E,
       }
     }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       if (!E || isa<ErrorExpr>(E) || !E->getType())
@@ -4236,7 +4264,11 @@ static void diagnoseDeprecatedWritableKeyPath(const Expr *E,
       }
     }
 
-    bool shouldWalkIntoNonSingleExpressionClosure() override { return false; }
+    bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+      return expr->hasAppliedFunctionBuilder();
+    }
+
+    bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       if (!E || isa<ErrorExpr>(E) || !E->getType())

lib/Sema/TypeCheckAvailability.cpp

@@ -2353,10 +2353,12 @@ class AvailabilityWalker : public ASTWalker {
     return true;
   }
 
-  bool shouldWalkIntoNonSingleExpressionClosure() override {
-    return false;
+  bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *expr) override {
+    return expr->hasAppliedFunctionBuilder();
   }
 
+  bool shouldWalkIntoTapExpression() override { return false; }
+
   std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
     ExprStack.push_back(E);
 

test/Constraints/function_builder_diags.swift

@@ -299,3 +299,31 @@ func checkSingleReturn(cond: Bool) {
     ($0, 17)
   }
 }
+
+// rdar://problem/59116520
+func checkImplicitSelfInClosure() {
+  @_functionBuilder
+  struct Builder {
+    static func buildBlock(_ children: String...) -> Element { Element() }
+  }
+
+  struct Element {
+    static func nonEscapingClosure(@Builder closure: (() -> Element)) {}
+    static func escapingClosure(@Builder closure: @escaping (() -> Element)) {}
+  }
+
+  class C {
+    let identifier: String = ""
+
+    func testImplicitSelf() {
+      Element.nonEscapingClosure {
+        identifier // okay
+      }
+
+      Element.escapingClosure { // expected-note {{capture 'self' explicitly to enable implicit 'self' in this closure}}
+        identifier // expected-error {{reference to property 'identifier' in closure requires explicit use of 'self' to make capture semantics explicit}}
+        // expected-note@-1 {{reference 'self.' explicitly}}
+      }
+    }
+  }
+}