Merge pull request swiftlang#32005 from DougGregor/function-builder-infer

DougGregor · web-flow · commit 8e0ec601db18 · 2020-05-26T10:59:42.000-07:00
[Function builders] Infer function builder from a protocol requirement.
diff --git a/include/swift/AST/ASTTypeIDZone.def b/include/swift/AST/ASTTypeIDZone.def
@@ -57,6 +57,7 @@ SWIFT_TYPEID_NAMED(PatternBindingEntry *, PatternBindingEntry)
 SWIFT_TYPEID_NAMED(PostfixOperatorDecl *, PostfixOperatorDecl)
 SWIFT_TYPEID_NAMED(PrecedenceGroupDecl *, PrecedenceGroupDecl)
 SWIFT_TYPEID_NAMED(PrefixOperatorDecl *, PrefixOperatorDecl)
+SWIFT_TYPEID_NAMED(ProtocolConformance *, ProtocolConformance)
 SWIFT_TYPEID_NAMED(ProtocolDecl *, ProtocolDecl)
 SWIFT_TYPEID_NAMED(SourceFile *, SourceFile)
 SWIFT_TYPEID_NAMED(TypeAliasDecl *, TypeAliasDecl)
diff --git a/include/swift/AST/ASTTypeIDs.h b/include/swift/AST/ASTTypeIDs.h
@@ -52,6 +52,7 @@ struct PropertyWrapperTypeInfo;
 enum class CtorInitializerKind;
 struct PropertyWrapperLValueness;
 struct PropertyWrapperMutability;
+class ProtocolConformance;
 class ProtocolDecl;
 class Requirement;
 enum class ResilienceExpansion : unsigned;
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -5069,6 +5069,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
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -2240,7 +2240,7 @@ class ClosureHasExplicitResultRequest
   bool isCached() const { return true; }
 };
 
-using ProtocolConformanceLookupResult = SmallVector<ProtocolConformance *, 2>;
+using ProtocolConformanceLookupResult = std::vector<ProtocolConformance *>;
 void simple_display(llvm::raw_ostream &out, ConformanceLookupKind kind);
 
 /// Lookup and expand all conformances in the given context.
@@ -2262,7 +2262,7 @@ class LookupAllConformancesInContextRequest
     : public SimpleRequest<LookupAllConformancesInContextRequest,
                            ProtocolConformanceLookupResult(
                                const IterableDeclContext *),
-                           RequestFlags::Uncached |
+                           RequestFlags::Cached |
                                RequestFlags::DependencySink |
                                RequestFlags::DependencySource> {
 public:
@@ -2276,6 +2276,8 @@ class LookupAllConformancesInContextRequest
   evaluate(Evaluator &evaluator, const IterableDeclContext *IDC) const;
 
 public:
+  bool isCached() const { return true; }
+
   // Incremental dependencies
   evaluator::DependencySource
   readDependencySource(const evaluator::DependencyRecorder &eval) const;
diff --git a/include/swift/Basic/SimpleDisplay.h b/include/swift/Basic/SimpleDisplay.h
@@ -136,6 +136,20 @@ namespace swift {
     out << "}";
   }
 
+  template<typename T>
+  void simple_display(llvm::raw_ostream &out,
+                      const std::vector<T> &vec) {
+    out << "{";
+    bool first = true;
+    for (const T &value : vec) {
+      if (first) first = false;
+      else out << ", ";
+
+      simple_display(out, value);
+    }
+    out << "}";
+  }
+
   template<typename T, typename U>
   void simple_display(llvm::raw_ostream &out,
                       const llvm::PointerUnion<T, U> &ptrUnion) {
diff --git a/lib/AST/Attr.cpp b/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;
       }
     }
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -6324,10 +6324,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,
@@ -6336,10 +6332,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,
diff --git a/lib/Sema/BuilderTransform.cpp b/lib/Sema/BuilderTransform.cpp
@@ -1401,10 +1401,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
@@ -1708,7 +1704,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();
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -2971,9 +2971,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;
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
@@ -5037,10 +5037,11 @@ diagnoseMissingAppendInterpolationMethod(NominalTypeDecl *typeDecl) {
   }
 }
 
-SmallVector<ProtocolConformance *, 2>
+std::vector<ProtocolConformance *>
 LookupAllConformancesInContextRequest::evaluate(
     Evaluator &eval, const IterableDeclContext *IDC) const {
-  return IDC->getLocalConformances(ConformanceLookupKind::All);
+  auto result = IDC->getLocalConformances(ConformanceLookupKind::All);
+  return std::vector<ProtocolConformance *>(result.begin(), result.end());
 }
 
 void TypeChecker::checkConformancesInContext(IterableDeclContext *idc) {
diff --git a/lib/Sema/TypeCheckRequestFunctions.cpp b/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,115 @@ 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() || isa<ProtocolDecl>(dc))
+    return Type();
+
+  auto funcDecl = dyn_cast<FuncDecl>(decl);
+  if (!funcDecl || !funcDecl->hasBody() ||
+      !decl->getDeclContext()->getParentSourceFile())
+    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();
+
+  // 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();
+
+  // 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);
diff --git a/lib/Sema/TypeCheckStmt.cpp b/lib/Sema/TypeCheckStmt.cpp
@@ -1862,7 +1862,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;
 
diff --git a/lib/Sema/TypeChecker.h b/lib/Sema/TypeChecker.h
@@ -561,6 +561,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);
diff --git a/test/Constraints/function_builder_infer.swift b/test/Constraints/function_builder_infer.swift
diff --git a/test/Incremental/Verifier/single-file-private/AnyObject.swift b/test/Incremental/Verifier/single-file-private/AnyObject.swift
diff --git a/test/Incremental/Verifier/single-file/AnyObject.swift b/test/Incremental/Verifier/single-file/AnyObject.swift
diff --git a/test/ModuleInterface/function_builders.swift b/test/ModuleInterface/function_builders.swift
diff --git a/test/Serialization/function_builders.swift b/test/Serialization/function_builders.swift

Original file line number	Diff line number	Diff line change
`@@ -5037,10 +5037,11 @@ diagnoseMissingAppendInterpolationMethod(NominalTypeDecl *typeDecl) {`
`5037`	`5037`	`}`
`5038`	`5038`	`}`
`5039`	`5039`
`5040`		`-SmallVector<ProtocolConformance *, 2>`
	`5040`	`+std::vector<ProtocolConformance *>`
`5041`	`5041`	`LookupAllConformancesInContextRequest::evaluate(`
`5042`	`5042`	`Evaluator &eval, const IterableDeclContext *IDC) const {`
`5043`		`- return IDC->getLocalConformances(ConformanceLookupKind::All);`
	`5043`	`+ auto result = IDC->getLocalConformances(ConformanceLookupKind::All);`
	`5044`	`+ return std::vector<ProtocolConformance *>(result.begin(), result.end());`
`5044`	`5045`	`}`
`5045`	`5046`
`5046`	`5047`	`void TypeChecker::checkConformancesInContext(IterableDeclContext *idc) {`