[Function builders] Infer function builder from a protocol requirement.

Allow a protocol requirement for a function or property to declare
a function builder. A witness to such a protocol requirement will
infer that function builder when all of the following are two:
* The witness does not explicitly adopt a function builder
* The witness is declared within the same context that conforms to the
protocol requirement (e.g., same extension or primary type declaration)
* The witness's body does not contain a "return" statement (since those
disable the function builder transform).

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -5048,6 +5048,14 @@ NOTE(function_builder_remove_attr, none,
      "remove the attribute to explicitly disable the function builder", ())
 NOTE(function_builder_remove_returns, none,
      "remove 'return' statements to apply the function builder", ())
+ERROR(function_builder_infer_ambig, none,
+      "ambiguous function builder inferred for %0: %1 or %2",
+      (DeclName, Type, Type))
+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))
 
 //------------------------------------------------------------------------------
 // MARK: Tuple Shuffle Diagnostics

lib/AST/Attr.cpp

@@ -758,11 +758,13 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
   case DAK_Custom: {
     if (!Options.IsForSwiftInterface)
       break;
-    // For Swift interface, we should only print function builder attribute
-    // on parameter decls. Printing the attribute elsewhere isn't ABI relevant.
+    // For Swift interface, we should print function builder attributes
+    // on parameter decls and on protocol requirements.
+    // Printing the attribute elsewhere isn't ABI relevant.
     if (auto *VD = dyn_cast<ValueDecl>(D)) {
       if (VD->getAttachedFunctionBuilder() == this) {
-        if (!isa<ParamDecl>(D))
+        if (!isa<ParamDecl>(D) &&
+            !(isa<VarDecl>(D) && isa<ProtocolDecl>(D->getDeclContext())))
           return false;
       }
     }

lib/AST/Decl.cpp

@@ -6345,10 +6345,6 @@ void ParamDecl::setStoredProperty(VarDecl *var) {
 }
 
 Type ValueDecl::getFunctionBuilderType() const {
-  // Fast path: most declarations (especially parameters, which is where
-  // this is hottest) do not have any custom attributes at all.
-  if (!getAttrs().hasAttribute<CustomAttr>()) return Type();
-
   auto &ctx = getASTContext();
   auto mutableThis = const_cast<ValueDecl *>(this);
   return evaluateOrDefault(ctx.evaluator,
@@ -6357,10 +6353,6 @@ Type ValueDecl::getFunctionBuilderType() const {
 }
 
 CustomAttr *ValueDecl::getAttachedFunctionBuilder() const {
-  // Fast path: most declarations (especially parameters, which is where
-  // this is hottest) do not have any custom attributes at all.
-  if (!getAttrs().hasAttribute<CustomAttr>()) return nullptr;
-
   auto &ctx = getASTContext();
   auto mutableThis = const_cast<ValueDecl *>(this);
   return evaluateOrDefault(ctx.evaluator,

lib/Sema/BuilderTransform.cpp

@@ -1223,10 +1223,6 @@ BraceStmt *swift::applyFunctionBuilderTransform(
         captured.first, captured.second)));
 }
 
-/// Find the return statements in the given body, which block the application
-/// of a function builder.
-static std::vector<ReturnStmt *> findReturnStatements(AnyFunctionRef fn);
-
 Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
     FuncDecl *func, Type builderType) {
   // Pre-check the body: pre-check any expressions in it and look
@@ -1552,7 +1548,7 @@ PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
   return PreCheckFunctionBuilderApplication(fn, false).run();
 }
 
-std::vector<ReturnStmt *> findReturnStatements(AnyFunctionRef fn) {
+std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
   PreCheckFunctionBuilderApplication precheck(fn, true);
   (void)precheck.run();
   return precheck.getReturnStmts();

lib/Sema/TypeCheckAttr.cpp

@@ -2941,9 +2941,11 @@ void AttributeChecker::visitCustomAttr(CustomAttr *attr) {
 
         // Module interfaces don't print bodies for all getters, so allow getters
         // that don't have a body if we're compiling a module interface.
+        // Within a protocol definition, there will never be a body.
         SourceFile *parent = storage->getDeclContext()->getParentSourceFile();
         bool isInInterface = parent && parent->Kind == SourceFileKind::Interface;
-        if (!isInInterface && !getter->hasBody())
+        if (!isInInterface && !getter->hasBody() &&
+            !isa<ProtocolDecl>(storage->getDeclContext()))
           return true;
 
         return false;

lib/Sema/TypeCheckRequestFunctions.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/NameLookupRequests.h"
+#include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeLoc.h"
 #include "swift/AST/Types.h"
 #include "swift/Subsystems.h"
@@ -182,11 +183,114 @@ AttachedFunctionBuilderRequest::evaluate(Evaluator &evaluator,
   return nullptr;
 }
 
+/// Attempt to infer the function builder type for a declaration.
+static Type inferFunctionBuilderType(ValueDecl *decl)  {
+  auto dc = decl->getDeclContext();
+  if (!dc->isTypeContext())
+    return Type();
+
+  auto funcDecl = dyn_cast<FuncDecl>(decl);
+  if (!funcDecl)
+    return Type();
+
+  // Only getters can have function builders. When we find one, look at
+  // the storage declaration for the purposes of witness matching.
+  auto lookupDecl = decl;
+  if (auto accessor = dyn_cast<AccessorDecl>(funcDecl)) {
+    if (accessor->getAccessorKind() != AccessorKind::Get)
+      return Type();
+
+    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;
+    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;
+
+      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(decl->getModuleContext());
+      Type subFunctionBuilderType = functionBuilderType.subst(subs);
+
+      matches.push_back({conformance, requirement, subFunctionBuilderType});
+    }
+  }
+
+  if (matches.size() == 0)
+    return Type();
+
+  // Check whether there are any return statements in the function's body.
+  // If there are, the function builder transform will be disabled,
+  // so don't infer a function builder.
+  if (!TypeChecker::findReturnStatements(funcDecl).empty())
+    return Type();
+
+  // Determine whether there is more than one actual function builder type.
+  Type functionBuilderType = matches[0].functionBuilderType;
+  for (const auto &match : matches) {
+    // If the types were the same anyway, there's nothing to do.
+    Type otherFunctionBuilderType = match.functionBuilderType;
+    if (functionBuilderType->isEqual(otherFunctionBuilderType))
+      continue;
+
+    // We have at least two different function builder types.
+    // Diagnose the ambiguity and provide potential solutions.
+    decl->diagnose(
+        diag::function_builder_infer_ambig, lookupDecl->getName(),
+        functionBuilderType, otherFunctionBuilderType);
+    decl->diagnose(diag::function_builder_infer_add_return)
+      .fixItInsert(funcDecl->getBodySourceRange().End, "return <#expr#>\n");
+    for (const auto &match : matches) {
+      decl->diagnose(
+          diag::function_builder_infer_pick_specific,
+          match.functionBuilderType,
+          match.conformance->getProtocol()->getName())
+        .fixItInsert(
+          lookupDecl->getAttributeInsertionLoc(false),
+          "@" + match.functionBuilderType.getString() + " ");
+    }
+
+    return Type();
+  }
+
+  return functionBuilderType;
+}
+
 Type FunctionBuilderTypeRequest::evaluate(Evaluator &evaluator,
                                           ValueDecl *decl) const {
   // Look for a function-builder custom attribute.
   auto attr = decl->getAttachedFunctionBuilder();
-  if (!attr) return Type();
+  if (!attr)
+    return inferFunctionBuilderType(decl);
 
   // Resolve a type for the attribute.
   auto mutableAttr = const_cast<CustomAttr*>(attr);

lib/Sema/TypeCheckStmt.cpp

@@ -1871,7 +1871,8 @@ TypeCheckFunctionBodyUntilRequest::evaluate(Evaluator &evaluator,
   if (auto *func = dyn_cast<FuncDecl>(AFD)) {
     if (Type builderType = getFunctionBuilderType(func)) {
       if (auto optBody =
-            TypeChecker::applyFunctionBuilderBodyTransform(func, builderType)) {
+              TypeChecker::applyFunctionBuilderBodyTransform(
+                func, builderType)) {
         if (!*optBody)
           return true;
 

lib/Sema/TypeChecker.h

@@ -601,6 +601,9 @@ bool typeCheckAbstractFunctionBody(AbstractFunctionDecl *AFD);
 Optional<BraceStmt *> applyFunctionBuilderBodyTransform(FuncDecl *func,
                                                         Type builderType);
 
+/// Find the return statements within the body of the given function.
+std::vector<ReturnStmt *> findReturnStatements(AnyFunctionRef fn);
+
 bool typeCheckClosureBody(ClosureExpr *closure);
 
 bool typeCheckTapBody(TapExpr *expr, DeclContext *DC);