[Diagnostics] Make sure that fixes associated with function builders are found and diagnosed

Current logic in `applySolutionFixes` didn't actually look inside
of single-expression closures (because of brace statement) or
closures which are transformed by function builders.

Resolves: rdar://problem/51167632

Author: xedin

lib/Sema/CSApply.cpp

@@ -7458,44 +7458,78 @@ Expr *ConstraintSystem::coerceToRValue(Expr *expr) {
 /// Emit the fixes computed as part of the solution, returning true if we were
 /// able to emit an error message, or false if none of the fixits worked out.
 bool ConstraintSystem::applySolutionFixes(Expr *E, const Solution &solution) {
-  llvm::SmallDenseMap<Expr *, SmallVector<const ConstraintFix *, 4>>
-      fixesPerExpr;
-
+  // First transfer all of the deduced information back
+  // to the constraint system.
   applySolution(solution);
 
-  for (auto *fix : solution.Fixes)
-    fixesPerExpr[fix->getAnchor()].push_back(fix);
+  class DiagnosticWalker : public ASTWalker {
+    Expr *root;
+    const Solution &solution;
+    llvm::SmallDenseMap<Expr *, SmallVector<ConstraintFix *, 4>> fixesPerExpr;
 
-  auto diagnoseExprFailures = [&](Expr *expr) -> bool {
-    auto fixes = fixesPerExpr.find(expr);
-    if (fixes == fixesPerExpr.end())
-      return false;
+    /// Determines whether any error have been diagnosed while
+    /// trying to apply fixes associated with a given solution.
+    bool DiagnosedAnyErrors = false;
+
+  public:
+    DiagnosticWalker(Expr *expr, const Solution &solution)
+        : root(expr), solution(solution) {
+      for (auto *fix : solution.Fixes)
+        fixesPerExpr[fix->getAnchor()].push_back(fix);
+    }
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+      // Diagnose root expression last.
+      if (E == root)
+        return {true, E};
+
+      if (auto *closure = dyn_cast<ClosureExpr>(E)) {
+        auto result = solution.builderTransformedClosures.find(closure);
+        if (result != solution.builderTransformedClosures.end()) {
+          auto *transformedExpr = result->second.second;
+          // Since this closure has been transformed into something
+          // else let's look inside transformed expression instead.
+          return {true, transformedExpr};
+        }
+      }
 
-    bool diagnosedError = false;
-    for (const auto *fix : fixes->second) {
-      auto diagnosed = fix->diagnose(E);
+      diagnose(E);
+      return {true, E};
+    }
 
-      if (fix->isWarning()) {
-        assert(diagnosed && "warnings should always be diagnosed");
-        (void)diagnosed;
-      } else {
-        diagnosedError |= diagnosed;
+    Expr *walkToExprPost(Expr *E) override {
+      if (E == root)
+        diagnose(E);
+      return E;
+    }
+
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+      return {true, S};
+    }
+
+    bool hadErrors() const { return DiagnosedAnyErrors; }
+
+  private:
+    void diagnose(Expr *E) {
+      auto fixes = fixesPerExpr.find(E);
+      if (fixes == fixesPerExpr.end())
+        return;
+
+      for (const auto *fix : fixes->second) {
+        auto diagnosed = fix->diagnose(root);
+        if (fix->isWarning()) {
+          assert(diagnosed && "warnings should always be diagnosed");
+          (void)diagnosed;
+        } else {
+          DiagnosedAnyErrors |= diagnosed;
+        }
       }
     }
-    return diagnosedError;
   };
 
-  bool diagnosedError = false;
-  E->forEachChildExpr([&](Expr *subExpr) -> Expr * {
-    // Diagnose root expression at the end to
-    // preserve ordering.
-    if (subExpr != E)
-      diagnosedError |= diagnoseExprFailures(subExpr);
-    return subExpr;
-  });
-
-  diagnosedError |= diagnoseExprFailures(E);
-  return diagnosedError;
+  DiagnosticWalker diagnostics(E, solution);
+  E->walk(diagnostics);
+  return diagnostics.hadErrors();
 }
 
 /// Apply a given solution to the expression, producing a fully

test/Constraints/function_builder_diags.swift

@@ -1,4 +1,4 @@
-// RUN: %target-typecheck-verify-swift
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -enable-opaque-result-types
 
 @_functionBuilder
 struct TupleBuilder { // expected-note 2{{struct 'TupleBuilder' declared here}}
@@ -126,3 +126,48 @@ func testOverloading(name: String) {
     }
   }
 }
+
+protocol P {
+  associatedtype T
+}
+
+struct AnyP : P {
+  typealias T = Any
+  init<T>(_: T) where T : P {}
+}
+
+struct TupleP<U> : P {
+  typealias T = U
+  init(_: U) {}
+}
+
+@_functionBuilder
+struct Builder {
+  static func buildBlock<S0, S1>(_ stmt1: S0, _ stmt2: S1) // expected-note {{where 'S1' = 'Label<Any>.Type'}}
+           -> TupleP<(S0, S1)> where S0: P, S1: P {
+    return TupleP((stmt1, stmt2))
+  }
+}
+
+struct G<C> : P where C : P {
+  typealias T = C
+  init(@Builder _: () -> C) {}
+}
+
+struct Text : P {
+  typealias T = String
+  init(_: T) {}
+}
+
+struct Label<L> : P where L : P { // expected-note {{'L' declared as parameter to type 'Label'}}
+  typealias T = L
+  init(@Builder _: () -> L) {}
+}
+
+func test_51167632() -> some P {
+  AnyP(G { // expected-error {{static method 'buildBlock' requires that 'Label<Any>.Type' conform to 'P'}}
+    Text("hello")
+    Label  // expected-error {{generic parameter 'L' could not be inferred}}
+    // expected-note@-1 {{explicitly specify the generic arguments to fix this issue}} {{10-10=<<#L: P#>>}}
+  })
+}