[CS] Use performSyntacticDiagnosticsForTarget for function builder

Add a case for `SolutionApplicationTarget::Kind::function`
and move the logic over.

Author: hamishknight

lib/Sema/BuilderTransform.cpp

@@ -1453,64 +1453,6 @@ BraceStmt *swift::applyFunctionBuilderTransform(
         captured.first, captured.second)));
 }
 
-/// Produce any additional syntactic diagnostics for the body of a function
-/// that had a function builder applied.
-static void performAddOnDiagnostics(BraceStmt *stmt, DeclContext *dc) {
-  class AddOnDiagnosticWalker : public ASTWalker {
-    SmallVector<DeclContext *, 4> dcStack;
-
-  public:
-    AddOnDiagnosticWalker(DeclContext *dc) {
-      dcStack.push_back(dc);
-    }
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      performSyntacticExprDiagnostics(
-          expr, dcStack.back(), /*isExprStmt=*/false);
-
-      if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-        if (closure->wasSeparatelyTypeChecked()) {
-          dcStack.push_back(closure);
-          return { true, expr };
-        }
-      }
-
-      return { false, expr };
-    }
-
-    Expr *walkToExprPost(Expr *expr) override {
-      if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-        if (closure->wasSeparatelyTypeChecked()) {
-          assert(dcStack.back() == closure);
-          dcStack.pop_back();
-        }
-      }
-
-      return expr;
-    }
-
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
-      performStmtDiagnostics(dcStack.back()->getASTContext(), stmt);
-      return { true, stmt };
-    }
-
-    std::pair<bool, Pattern*> walkToPatternPre(Pattern *pattern) override {
-      return { false, pattern };
-    }
-
-    bool walkToTypeLocPre(TypeLoc &typeLoc) override { return false; }
-
-    bool walkToTypeReprPre(TypeRepr *typeRepr) override { return false; }
-
-    bool walkToParameterListPre(ParameterList *params) override {
-      return false;
-    }
-  };
-
-  AddOnDiagnosticWalker walker(dc);
-  stmt->walk(walker);
-}
-
 Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
     FuncDecl *func, Type builderType) {
   // Pre-check the body: pre-check any expressions in it and look
@@ -1629,7 +1571,7 @@ Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
   if (auto result = cs.applySolution(
           solutions.front(),
           SolutionApplicationTarget(func))) {
-    performAddOnDiagnostics(result->getFunctionBody(), func);
+    performSyntacticDiagnosticsForTarget(*result, /*isExprStmt*/ false);
     return result->getFunctionBody();
   }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -1995,6 +1995,54 @@ bool GenericRequirementsCheckListener::diagnoseUnsatisfiedRequirement(
   return false;
 }
 
+namespace {
+/// Produce any additional syntactic diagnostics for the body of a function
+/// that had a function builder applied.
+class FunctionSyntacticDiagnosticWalker : public ASTWalker {
+  SmallVector<DeclContext *, 4> dcStack;
+
+public:
+  FunctionSyntacticDiagnosticWalker(DeclContext *dc) { dcStack.push_back(dc); }
+
+  std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+    performSyntacticExprDiagnostics(expr, dcStack.back(), /*isExprStmt=*/false);
+
+    if (auto closure = dyn_cast<ClosureExpr>(expr)) {
+      if (closure->wasSeparatelyTypeChecked()) {
+        dcStack.push_back(closure);
+        return {true, expr};
+      }
+    }
+
+    return {false, expr};
+  }
+
+  Expr *walkToExprPost(Expr *expr) override {
+    if (auto closure = dyn_cast<ClosureExpr>(expr)) {
+      if (closure->wasSeparatelyTypeChecked()) {
+        assert(dcStack.back() == closure);
+        dcStack.pop_back();
+      }
+    }
+
+    return expr;
+  }
+
+  std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
+    performStmtDiagnostics(dcStack.back()->getASTContext(), stmt);
+    return {true, stmt};
+  }
+
+  std::pair<bool, Pattern *> walkToPatternPre(Pattern *pattern) override {
+    return {false, pattern};
+  }
+
+  bool walkToTypeLocPre(TypeLoc &typeLoc) override { return false; }
+  bool walkToTypeReprPre(TypeRepr *typeRepr) override { return false; }
+  bool walkToParameterListPre(ParameterList *params) override { return false; }
+};
+} // end anonymous namespace
+
 void constraints::performSyntacticDiagnosticsForTarget(
     const SolutionApplicationTarget &target, bool isExprStmt) {
   auto *dc = target.getDeclContext();
@@ -2011,7 +2059,11 @@ void constraints::performSyntacticDiagnosticsForTarget(
     }
     return;
   }
-  case SolutionApplicationTarget::Kind::function:
+  case SolutionApplicationTarget::Kind::function: {
+    FunctionSyntacticDiagnosticWalker walker(dc);
+    target.getFunctionBody()->walk(walker);
+    return;
+  }
   case SolutionApplicationTarget::Kind::stmtCondition:
   case SolutionApplicationTarget::Kind::caseLabelItem:
   case SolutionApplicationTarget::Kind::patternBinding: