[Function builders] Infer function builder through @_dynamicReplacement(for:)

include/swift/AST/DiagnosticsSema.def

@@ -5075,8 +5075,8 @@ ERROR(function_builder_infer_ambig, none,
 NOTE(function_builder_infer_add_return, none,
      "add an explicit 'return' statement to not use a function builder", ())
 NOTE(function_builder_infer_pick_specific, none,
-     "apply function builder %0 (inferred from protocol %1)",
-     (Type, DeclName))
+     "apply function builder %0 (inferred from %select{protocol|dynamic replacement of}1 %2)",
+     (Type, unsigned, DeclName))
 
 //------------------------------------------------------------------------------
 // MARK: Tuple Shuffle Diagnostics

lib/Sema/TypeCheckRequestFunctions.cpp

@@ -210,45 +210,109 @@ static Type inferFunctionBuilderType(ValueDecl *decl)  {
     lookupDecl = accessor->getStorage();
   }
 
-  // Determine all of the conformances within the same context as
-  // this declaration. If this declaration is a witness to any
-  // requirement within one of those protocols that has a function builder
-  // attached, use that function builder type.
-  auto idc = cast<IterableDeclContext>(dc->getAsDecl());
-  auto conformances = evaluateOrDefault(
-      dc->getASTContext().evaluator,
-      LookupAllConformancesInContextRequest{idc}, { });
-
   // Find all of the potentially inferred function builder types.
   struct Match {
-    ProtocolConformance *conformance;
-    ValueDecl *requirement;
+    enum Kind {
+      Conformance,
+      DynamicReplacement,
+    } kind;
+
+    union {
+      struct {
+        ProtocolConformance *conformance;
+        ValueDecl *requirement;
+      } conformanceMatch;
+
+      ValueDecl *dynamicReplacement;
+    };
+
     Type functionBuilderType;
-  };
-  SmallVector<Match, 2> matches;
-  for (auto conformance : conformances) {
-    auto protocol = conformance->getProtocol();
-    for (auto found : protocol->lookupDirect(lookupDecl->getName())) {
-      if (!isa<ProtocolDecl>(found->getDeclContext()))
-        continue;
 
-      auto requirement = dyn_cast<ValueDecl>(found);
-      if (!requirement)
-        continue;
+    static Match forConformance(
+        ProtocolConformance *conformance,
+        ValueDecl *requirement,
+        Type functionBuilderType) {
+      Match match;
+      match.kind = Conformance;
+      match.conformanceMatch.conformance = conformance;
+      match.conformanceMatch.requirement = requirement;
+      match.functionBuilderType = functionBuilderType;
+      return match;
+    }
 
-      Type functionBuilderType = requirement->getFunctionBuilderType();
-      if (!functionBuilderType)
-        continue;
+    static Match forDynamicReplacement(
+        ValueDecl *dynamicReplacement, Type functionBuilderType) {
+      Match match;
+      match.kind = DynamicReplacement;
+      match.dynamicReplacement = dynamicReplacement;
+      match.functionBuilderType = functionBuilderType;
+      return match;
+    }
 
-      auto witness = conformance->getWitnessDecl(requirement);
-      if (witness != lookupDecl)
-        continue;
+    DeclName getSourceName() const {
+      switch (kind) {
+      case Conformance:
+        return conformanceMatch.conformance->getProtocol()->getName();
 
-      // Substitute into the function builder type.
-      auto subs = conformance->getSubstitutions(decl->getModuleContext());
-      Type subFunctionBuilderType = functionBuilderType.subst(subs);
+      case DynamicReplacement:
+        return dynamicReplacement->getName();
+      }
+    }
+  };
+
+  // The set of matches from which we can infer function builder types.
+  SmallVector<Match, 2> matches;
 
-      matches.push_back({conformance, requirement, subFunctionBuilderType});
+  // Determine all of the conformances within the same context as
+  // this declaration. If this declaration is a witness to any
+  // requirement within one of those protocols that has a function builder
+  // attached, use that function builder type.
+  auto addConformanceMatches = [&matches](ValueDecl *lookupDecl) {
+    DeclContext *dc = lookupDecl->getDeclContext();
+    auto idc = cast<IterableDeclContext>(dc->getAsDecl());
+    auto conformances = evaluateOrDefault(
+        dc->getASTContext().evaluator,
+        LookupAllConformancesInContextRequest{idc}, { });
+
+    for (auto conformance : conformances) {
+      auto protocol = conformance->getProtocol();
+      for (auto found : protocol->lookupDirect(lookupDecl->getName())) {
+        if (!isa<ProtocolDecl>(found->getDeclContext()))
+          continue;
+
+        auto requirement = dyn_cast<ValueDecl>(found);
+        if (!requirement)
+          continue;
+
+        Type functionBuilderType = requirement->getFunctionBuilderType();
+        if (!functionBuilderType)
+          continue;
+
+        auto witness = conformance->getWitnessDecl(requirement);
+        if (witness != lookupDecl)
+          continue;
+
+        // Substitute into the function builder type.
+        auto subs =
+            conformance->getSubstitutions(lookupDecl->getModuleContext());
+        Type subFunctionBuilderType = functionBuilderType.subst(subs);
+
+        matches.push_back(
+            Match::forConformance(
+              conformance, requirement, subFunctionBuilderType));
+      }
+    }
+  };
+
+  addConformanceMatches(lookupDecl);
+
+  // Look for function builder types inferred through dynamic replacements.
+  if (auto replaced = lookupDecl->getDynamicallyReplacedDecl()) {
+    if (auto functionBuilderType = replaced->getFunctionBuilderType()) {
+      matches.push_back(
+        Match::forDynamicReplacement(replaced, functionBuilderType));
+    } else {
+      addConformanceMatches(replaced);
     }
   }
 
@@ -274,7 +338,8 @@ static Type inferFunctionBuilderType(ValueDecl *decl)  {
       decl->diagnose(
           diag::function_builder_infer_pick_specific,
           match.functionBuilderType,
-          match.conformance->getProtocol()->getName())
+          static_cast<unsigned>(match.kind),
+          match.getSourceName())
         .fixItInsert(
           lookupDecl->getAttributeInsertionLoc(false),
           "@" + match.functionBuilderType.getString() + " ");

test/Constraints/function_builder_infer.swift

@@ -184,3 +184,51 @@ struct TupleMeResolvedExplicit: Tupled, OtherTupled {
     return "hello"
   }
 }
+
+// Inference through dynamic replacement
+struct DynamicTupled: Tupled {
+  dynamic var tuple: some Any {
+    return "hello"
+  }
+}
+
+extension DynamicTupled {
+  @_dynamicReplacement(for: tuple)
+  var replacementTuple: some Any {
+    1
+    3.14159
+    "hello"
+  }
+}
+
+struct DynamicTupled2: Tupled, OtherTupled {
+  dynamic var tuple: some Any {
+    return "hello"
+  }
+}
+
+extension DynamicTupled2 {
+  @_dynamicReplacement(for: tuple)
+  var replacementTuple: some Any { // expected-error{{ambiguous function builder inferred for 'replacementTuple': 'TupleBuilder' or 'OtherTupleBuilder'}}
+    // expected-note@-1{{add an explicit 'return' statement to not use a function builder}}
+    // expected-note@-2{{apply function builder 'TupleBuilder' (inferred from protocol 'Tupled')}}
+    // expected-note@-3{{apply function builder 'OtherTupleBuilder' (inferred from protocol 'OtherTupled')}}
+    1
+  }
+}
+
+struct DynamicTupled3 {
+  @TupleBuilder dynamic var dynamicTuple: some Any {
+    0
+  }
+}
+
+extension DynamicTupled3: OtherTupled {
+  @_dynamicReplacement(for: dynamicTuple)
+  var tuple: some Any { // expected-error{{ambiguous function builder inferred for 'tuple': 'OtherTupleBuilder' or 'TupleBuilder'}}
+    // expected-note@-1{{add an explicit 'return' statement to not use a function builder}}
+    // expected-note@-2{{apply function builder 'OtherTupleBuilder' (inferred from protocol 'OtherTupled')}}
+    // expected-note@-3{{apply function builder 'TupleBuilder' (inferred from dynamic replacement of 'dynamicTuple')}}
+    0
+  }
+}