[ConstraintSystem] Handle requirement failures associated with injected result builder methods

Only way to aggregate build* methods is via source locations because
each result builder kind would create a new build* expression node.

(cherry picked from commit e99f0e2)

Author: xedin

lib/Sema/ConstraintSystem.cpp

@@ -4860,66 +4860,68 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
   // Aggregate all requirement fixes that belong to the same callee
   // and attempt to diagnose possible ambiguities.
   {
-    auto isResultBuilderRef = [&](ASTNode node) {
+    auto isResultBuilderMethodRef = [&](ASTNode node) {
       auto *UDE = getAsExpr<UnresolvedDotExpr>(node);
-      if (!UDE)
+      if (!(UDE && UDE->isImplicit()))
         return false;
 
       auto &ctx = getASTContext();
-      return UDE->isImplicit() &&
-        UDE->getName().compare(DeclNameRef(ctx.Id_buildBlock)) == 0;
+      SmallVector<Identifier, 4> builderMethods(
+          {ctx.Id_buildBlock, ctx.Id_buildExpression, ctx.Id_buildPartialBlock,
+           ctx.Id_buildFinalResult});
+
+      return llvm::any_of(builderMethods, [&](const Identifier &methodId) {
+        return UDE->getName().compare(DeclNameRef(methodId)) == 0;
+      });
     };
 
-    llvm::MapVector<SourceLoc,
-                    SmallVector<FixInContext, 4>>
-        builderFixes;
+    // Aggregates fixes fixes attached to `buildExpression` and `buildBlock`
+    // methods at the particular source location.
+    llvm::MapVector<SourceLoc, SmallVector<FixInContext, 4>>
+        builderMethodRequirementFixes;
 
     llvm::MapVector<ConstraintLocator *, SmallVector<FixInContext, 4>>
-        requirementFixes;
+        perCalleeRequirementFixes;
 
-    // TODO(diagnostics): This approach doesn't work for synthesized code
-    // because i.e. result builders would inject a new `buildBlock` or
-    // `buildExpression` for every kind of builder. We need to come up
-    // with the strategy to unify locators in such cases.
     for (const auto &entry : fixes) {
       auto *fix = entry.second;
       if (!fix->getLocator()->isLastElement<LocatorPathElt::AnyRequirement>())
         continue;
 
       auto *calleeLoc = entry.first->getCalleeLocator(fix->getLocator());
 
-      if (isResultBuilderRef(calleeLoc->getAnchor())) {
+      if (isResultBuilderMethodRef(calleeLoc->getAnchor())) {
         auto *anchor = castToExpr<Expr>(calleeLoc->getAnchor());
-        builderFixes[anchor->getLoc()].push_back(entry);
+        builderMethodRequirementFixes[anchor->getLoc()].push_back(entry);
       } else {
-        requirementFixes[calleeLoc].push_back(entry);
+        perCalleeRequirementFixes[calleeLoc].push_back(entry);
       }
     }
 
-    SmallVector<FixInContext, 4> diagnosedFixes;
+    SmallVector<SmallVector<FixInContext, 4>, 4> viableGroups;
+    {
+      auto takeAggregateIfViable =
+          [&](SmallVector<FixInContext, 4> &aggregate) {
+            // Ambiguity only if all of the solutions have a requirement
+            // fix at the given location.
+            if (aggregate.size() == solutions.size())
+              viableGroups.push_back(std::move(aggregate));
+          };
 
-    for (auto &entry : builderFixes) {
-      auto &aggregate = entry.second;
+      for (auto &entry : builderMethodRequirementFixes)
+        takeAggregateIfViable(entry.second);
 
-      if (diagnoseAmbiguityWithGenericRequirements(*this, aggregate))
-        diagnosedFixes.append(aggregate.begin(), aggregate.end());
+      for (auto &entry : perCalleeRequirementFixes)
+        takeAggregateIfViable(entry.second);
     }
 
-    for (auto &entry : requirementFixes) {
-      auto &aggregate = entry.second;
-
-      // Ambiguity only if all of the solutions have a requirement
-      // fix at the given location.
-      if (aggregate.size() != solutions.size())
-        continue;
-
-      if (diagnoseAmbiguityWithGenericRequirements(*this, aggregate))
-        diagnosedFixes.append(aggregate.begin(), aggregate.end());
+    for (auto &aggregate : viableGroups) {
+      if (diagnoseAmbiguityWithGenericRequirements(*this, aggregate)) {
+        // Remove diagnosed fixes.
+        fixes.set_subtract(aggregate);
+        diagnosed = true;
+      }
     }
-
-    diagnosed |= !diagnosedFixes.empty();
-    // Remove any diagnosed fixes.
-    fixes.set_subtract(diagnosedFixes);
   }
 
   llvm::MapVector<std::pair<FixKind, ConstraintLocator *>,

test/Constraints/result_builder_diags.swift

@@ -967,3 +967,36 @@ func test_impact_of_control_flow_fix() {
     test("") // expected-error {{cannot convert value of type 'String' to expected argument type 'Int'}}
   }
 }
+
+protocol Q {}
+
+func test_requirement_failure_in_buildBlock() {
+  struct A : P { typealias T = Int }
+  struct B : Q { typealias T = String }
+
+  struct Result<T> : P { typealias T = Int }
+
+  @resultBuilder
+  struct BuilderA {
+    static func buildBlock<T0: P, T1: P>(_: T0, _: T1) -> Result<(T0, T1)> { fatalError() }
+    // expected-note@-1 {{candidate requires that 'B' conform to 'P' (requirement specified as 'T1' : 'P')}}
+  }
+
+  @resultBuilder
+  struct BuilderB {
+    static func buildBlock<U0: Q, U1: Q>(_ v: U0, _: U1) -> some Q { v }
+    // expected-note@-1 {{candidate requires that 'A' conform to 'Q' (requirement specified as 'U0' : 'Q')}}
+  }
+
+  struct Test {
+    func fn<T: P>(@BuilderA _: () -> T) {}
+    func fn<T: Q>(@BuilderB _: () -> T) {}
+
+    func test() {
+      fn { // expected-error {{no exact matches in reference to static method 'buildBlock'}}
+        A()
+        B()
+      }
+    }
+  }
+}