[Constraint system] Generalize function builder application APIs.

Teach the constraint system to handle matching a function builder to a
function as well as a closure.

Author: DougGregor

include/swift/AST/AnyFunctionRef.h

@@ -216,6 +216,15 @@ class AnyFunctionRef {
      return lhs.TheFunction != rhs.TheFunction;
    }
 
+  friend llvm::hash_code hash_value(AnyFunctionRef fn) {
+    using llvm::hash_value;
+    return hash_value(fn.TheFunction.getOpaqueValue());
+  }
+
+  friend SourceLoc extractNearestSourceLoc(AnyFunctionRef fn) {
+    return fn.getLoc();
+  }
+
 private:
   ArrayRef<AnyFunctionType::Yield>
   getYieldResultsImpl(SmallVectorImpl<AnyFunctionType::Yield> &buffer,
@@ -243,6 +252,8 @@ class AnyFunctionRef {
 #pragma warning(pop)
 #endif
 
+void simple_display(llvm::raw_ostream &out, AnyFunctionRef fn);
+
 } // namespace swift
 
 namespace llvm {

include/swift/AST/TypeCheckRequests.h

@@ -16,6 +16,7 @@
 #ifndef SWIFT_TYPE_CHECK_REQUESTS_H
 #define SWIFT_TYPE_CHECK_REQUESTS_H
 
+#include "swift/AST/AnyFunctionRef.h"
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Type.h"
@@ -1760,7 +1761,7 @@ enum class FunctionBuilderClosurePreCheck : uint8_t {
 
 class PreCheckFunctionBuilderRequest
     : public SimpleRequest<PreCheckFunctionBuilderRequest,
-                           FunctionBuilderClosurePreCheck(ClosureExpr *),
+                           FunctionBuilderClosurePreCheck(AnyFunctionRef),
                            CacheKind::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -1770,7 +1771,7 @@ class PreCheckFunctionBuilderRequest
 
   // Evaluation.
   llvm::Expected<FunctionBuilderClosurePreCheck>
-  evaluate(Evaluator &evaluator, ClosureExpr *closure) const;
+  evaluate(Evaluator &evaluator, AnyFunctionRef fn) const;
 
 public:
   // Separate caching.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -196,7 +196,7 @@ SWIFT_REQUEST(TypeChecker, HasUserDefinedDesignatedInitRequest,
 SWIFT_REQUEST(TypeChecker, HasMemberwiseInitRequest,
               bool(StructDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, PreCheckFunctionBuilderRequest,
-              FunctionBuilderClosurePreCheck(ClosureExpr *),
+              FunctionBuilderClosurePreCheck(AnyFunctionRef),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ResolveImplicitMemberRequest,
               bool(NominalTypeDecl *, ImplicitMemberAction), Uncached,

lib/AST/Decl.cpp

@@ -7907,3 +7907,10 @@ void ParseAbstractFunctionBodyRequest::cacheResult(BraceStmt *value) const {
   }
 
 }
+
+void swift::simple_display(llvm::raw_ostream &out, AnyFunctionRef fn) {
+  if (auto func = fn.getAbstractFunctionDecl())
+    simple_display(out, func);
+  else
+    out << "closure";
+}

lib/Sema/BuilderTransform.cpp

@@ -510,21 +510,23 @@ TypeChecker::applyFunctionBuilderBodyTransform(FuncDecl *FD,
                            body->getRBraceLoc());
 }
 
-ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
-    ClosureExpr *closure, Type builderType, ConstraintLocator *calleeLocator,
-    ConstraintLocatorBuilder locator) {
+ConstraintSystem::TypeMatchResult ConstraintSystem::matchFunctionBuilder(
+    AnyFunctionRef fn, Type builderType, Type bodyResultType,
+    ConstraintLocator *calleeLocator, ConstraintLocatorBuilder locator) {
   auto builder = builderType->getAnyNominal();
   assert(builder && "Bad function builder type");
   assert(builder->getAttrs().hasAttribute<FunctionBuilderAttr>());
 
   // FIXME: Right now, single-expression closures suppress the function
   // builder translation.
-  if (closure->hasSingleExpressionBody())
-    return getTypeMatchSuccess();
+  if (auto closure = fn.getAbstractClosureExpr()) {
+    if (closure->hasSingleExpressionBody())
+      return getTypeMatchSuccess();
+  }
 
   // Pre-check the closure body: pre-check any expressions in it and look
   // for return statements.
-  auto request = PreCheckFunctionBuilderRequest{closure};
+  auto request = PreCheckFunctionBuilderRequest{fn};
   switch (evaluateOrDefault(getASTContext().evaluator, request,
                             FunctionBuilderClosurePreCheck::Error)) {
   case FunctionBuilderClosurePreCheck::Okay:
@@ -547,7 +549,7 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
     // Check whether we can apply this specific function builder.
     BuilderClosureVisitor visitor(getASTContext(), this,
                                   /*wantExpr=*/false, builderType);
-    (void)visitor.visit(closure->getBody());
+    (void)visitor.visit(fn.getBody());
 
     // If we saw a control-flow statement or declaration that the builder
     // cannot handle, we don't have a well-formed function builder application.
@@ -585,17 +587,18 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
 
   BuilderClosureVisitor visitor(getASTContext(), this,
                                 /*wantExpr=*/true, builderType);
-  Expr *singleExpr = visitor.visit(closure->getBody());
+  Expr *singleExpr = visitor.visit(fn.getBody());
 
   // We've already pre-checked all the original expressions, but do the
   // pre-check to the generated expression just to set up any preconditions
   // that CSGen might have.
   //
   // TODO: just build the AST the way we want it in the first place.
-  if (ConstraintSystem::preCheckExpression(singleExpr, closure))
+  auto dc = fn.getAsDeclContext();
+  if (ConstraintSystem::preCheckExpression(singleExpr, dc))
     return getTypeMatchFailure(locator);
 
-  singleExpr = generateConstraints(singleExpr, closure);
+  singleExpr = generateConstraints(singleExpr, dc);
   if (!singleExpr)
     return getTypeMatchFailure(locator);
 
@@ -607,41 +610,37 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
       functionBuilderTransformed.begin(),
       functionBuilderTransformed.end(),
       [&](const std::pair<AnyFunctionRef, AppliedBuilderTransform> &elt) {
-        return elt.first == closure;
+        return elt.first == fn;
       }) == functionBuilderTransformed.end() &&
          "already transformed this closure along this path!?!");
   functionBuilderTransformed.push_back(
-      std::make_pair(closure,
+      std::make_pair(fn,
                      AppliedBuilderTransform{builderType, singleExpr}));
 
-  // Bind the result type of the closure to the type of the transformed
-  // expression.
-  Type closureType = getType(closure);
-  auto fnType = closureType->castTo<FunctionType>();
-  addConstraint(ConstraintKind::Equal, fnType->getResult(), transformedType,
+  // Bind the body result type to the type of the transformed expression.
+  addConstraint(ConstraintKind::Equal, bodyResultType, transformedType,
                 locator);
   return getTypeMatchSuccess();
 }
 
 namespace {
 
 /// Pre-check all the expressions in the closure body.
-class PreCheckFunctionBuilderClosure : public ASTWalker {
-  ClosureExpr *Closure;
+class PreCheckFunctionBuilderApplication : public ASTWalker {
+  AnyFunctionRef Fn;
   bool HasReturnStmt = false;
   bool HasError = false;
 public:
-  PreCheckFunctionBuilderClosure(ClosureExpr *closure)
-    : Closure(closure) {}
+  PreCheckFunctionBuilderApplication(AnyFunctionRef fn) : Fn(fn) {}
 
   FunctionBuilderClosurePreCheck run() {
-    Stmt *oldBody = Closure->getBody();
+    Stmt *oldBody = Fn.getBody();
 
     Stmt *newBody = oldBody->walk(*this);
 
     // If the walk was aborted, it was because we had a problem of some kind.
     assert((newBody == nullptr) == (HasError || HasReturnStmt) &&
-           "unexpected short-circuit while walking closure body");
+           "unexpected short-circuit while walking body");
     if (!newBody) {
       if (HasError)
         return FunctionBuilderClosurePreCheck::Error;
@@ -658,7 +657,7 @@ class PreCheckFunctionBuilderClosure : public ASTWalker {
     // Pre-check the expression.  If this fails, abort the walk immediately.
     // Otherwise, replace the expression with the result of pre-checking.
     // In either case, don't recurse into the expression.
-    if (ConstraintSystem::preCheckExpression(E, /*DC*/ Closure)) {
+    if (ConstraintSystem::preCheckExpression(E, /*DC*/ Fn.getAsDeclContext())) {
       HasError = true;
       return std::make_pair(false, nullptr);
     }
@@ -682,10 +681,12 @@ class PreCheckFunctionBuilderClosure : public ASTWalker {
 
 llvm::Expected<FunctionBuilderClosurePreCheck>
 PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
-                                         ClosureExpr *closure) const {
+                                         AnyFunctionRef fn) const {
   // Single-expression closures should already have been pre-checked.
-  if (closure->hasSingleExpressionBody())
-    return FunctionBuilderClosurePreCheck::Okay;
+  if (auto closure = fn.getAbstractClosureExpr()) {
+    if (closure->hasSingleExpressionBody())
+      return FunctionBuilderClosurePreCheck::Okay;
+  }
 
-  return PreCheckFunctionBuilderClosure(closure).run();
+  return PreCheckFunctionBuilderApplication(fn).run();
 }

lib/Sema/CSSimplify.cpp

@@ -1145,9 +1145,11 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
               = paramInfo.getFunctionBuilderType(paramIdx)) {
         Expr *arg = getArgumentExpr(locator.getAnchor(), argIdx);
         if (auto closure = dyn_cast_or_null<ClosureExpr>(arg)) {
-          auto result =
-              cs.applyFunctionBuilder(closure, functionBuilderType,
-                                      calleeLocator, loc);
+          auto closureType = cs.getType(closure);
+          auto result = cs.matchFunctionBuilder(
+              closure, functionBuilderType,
+              closureType->castTo<FunctionType>()->getResult(),
+              calleeLocator, loc);
           if (result.isFailure())
             return result;
         }

lib/Sema/ConstraintSystem.h

@@ -3477,9 +3477,9 @@ class ConstraintSystem {
   void simplifyDisjunctionChoice(Constraint *choice);
 
   /// Apply the given function builder to the closure expression.
-  TypeMatchResult applyFunctionBuilder(ClosureExpr *closure, Type builderType,
-                                       ConstraintLocator *calleeLocator,
-                                       ConstraintLocatorBuilder locator);
+  TypeMatchResult matchFunctionBuilder(
+      AnyFunctionRef fn, Type builderType, Type bodyResultType,
+      ConstraintLocator *calleeLocator, ConstraintLocatorBuilder locator);
 
 private:
   /// The kind of bindings that are permitted.