Merge pull request #41189 from xedin/trailing-closures-with-callAsFunction

[ConstraintSystem] Match trailing closures to implicit `.callAsFunction` when necessary

Author: xedin

include/swift/AST/ArgumentList.h

@@ -249,11 +249,13 @@ class alignas(Argument) ArgumentList final
   static ArgumentList *
   createImplicit(ASTContext &ctx, SourceLoc lParenLoc, ArrayRef<Argument> args,
                  SourceLoc rParenLoc,
+                 Optional<unsigned> firstTrailingClosureIndex = None,
                  AllocationArena arena = AllocationArena::Permanent);
 
   /// Create a new implicit ArgumentList with a set of \p args.
   static ArgumentList *
   createImplicit(ASTContext &ctx, ArrayRef<Argument> args,
+                 Optional<unsigned> firstTrailingClosureIndex = None,
                  AllocationArena arena = AllocationArena::Permanent);
 
   /// Create a new implicit ArgumentList with a single labeled argument

include/swift/Sema/ConstraintSystem.h

@@ -1220,6 +1220,10 @@ class Solution {
   /// constructions) to the argument lists for the call to that locator.
   llvm::MapVector<ConstraintLocator *, ArgumentList *> argumentLists;
 
+  /// The set of implicitly generated `.callAsFunction` root expressions.
+  llvm::DenseMap<ConstraintLocator *, UnresolvedDotExpr *>
+      ImplicitCallAsFunctionRoots;
+
   /// Record a new argument matching choice for given locator that maps a
   /// single argument to a single parameter.
   void recordSingleArgMatchingChoice(ConstraintLocator *locator);
@@ -2468,6 +2472,12 @@ class ConstraintSystem {
   /// types.
   llvm::DenseMap<CanType, DynamicCallableMethods> DynamicCallableCache;
 
+  /// A cache of implicitly generated dot-member expressions used as roots
+  /// for some `.callAsFunction` calls. The key here is "base" locator for
+  /// the `.callAsFunction` member reference.
+  llvm::SmallMapVector<ConstraintLocator *, UnresolvedDotExpr *, 2>
+      ImplicitCallAsFunctionRoots;
+
 private:
   /// Describe the candidate expression for partial solving.
   /// This class used by shrink & solve methods which apply
@@ -2951,6 +2961,9 @@ class ConstraintSystem {
     /// The length of \c ArgumentLists.
     unsigned numArgumentLists;
 
+    /// The length of \c ImplicitCallAsFunctionRoots.
+    unsigned numImplicitCallAsFunctionRoots;
+
     /// The previous score.
     Score PreviousScore;
 
@@ -3475,6 +3488,11 @@ class ConstraintSystem {
   void recordMatchCallArgumentResult(ConstraintLocator *locator,
                                      MatchCallArgumentResult result);
 
+  /// Record implicitly generated `callAsFunction` with root at the
+  /// given expression, located at \c locator.
+  void recordCallAsFunction(UnresolvedDotExpr *root, ArgumentList *arguments,
+                            ConstraintLocator *locator);
+
   /// Walk a closure AST to determine its effects.
   ///
   /// \returns a function's extended info describing the effects, as
@@ -5474,6 +5492,7 @@ matchCallArguments(
 ConstraintSystem::TypeMatchResult
 matchCallArguments(ConstraintSystem &cs,
                    FunctionType *contextualType,
+                   ArgumentList *argumentList,
                    ArrayRef<AnyFunctionType::Param> args,
                    ArrayRef<AnyFunctionType::Param> params,
                    ConstraintKind subKind,

lib/AST/ArgumentList.cpp

@@ -108,19 +108,22 @@ ArgumentList *ArgumentList::createTypeChecked(ASTContext &ctx,
                 originalArgs->isImplicit(), originalArgs);
 }
 
-ArgumentList *ArgumentList::createImplicit(ASTContext &ctx, SourceLoc lParenLoc,
-                                           ArrayRef<Argument> args,
-                                           SourceLoc rParenLoc,
-                                           AllocationArena arena) {
-  return create(ctx, lParenLoc, args, rParenLoc,
-                /*firstTrailingClosureIdx*/ None, /*implicit*/ true,
+ArgumentList *
+ArgumentList::createImplicit(ASTContext &ctx, SourceLoc lParenLoc,
+                             ArrayRef<Argument> args, SourceLoc rParenLoc,
+                             Optional<unsigned> firstTrailingClosureIndex,
+                             AllocationArena arena) {
+  return create(ctx, lParenLoc, args, rParenLoc, firstTrailingClosureIndex,
+                /*implicit*/ true,
                 /*originalArgs*/ nullptr, arena);
 }
 
-ArgumentList *ArgumentList::createImplicit(ASTContext &ctx,
-                                           ArrayRef<Argument> args,
-                                           AllocationArena arena) {
-  return createImplicit(ctx, SourceLoc(), args, SourceLoc(), arena);
+ArgumentList *
+ArgumentList::createImplicit(ASTContext &ctx, ArrayRef<Argument> args,
+                             Optional<unsigned> firstTrailingClosureIndex,
+                             AllocationArena arena) {
+  return createImplicit(ctx, SourceLoc(), args, SourceLoc(),
+                        firstTrailingClosureIndex, arena);
 }
 
 ArgumentList *ArgumentList::forImplicitSingle(ASTContext &ctx, Identifier label,

lib/Sema/CSApply.cpp

@@ -7712,7 +7712,38 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
   apply->setFn(declRef);
 
   // Tail-recur to actually call the constructor.
-  return finishApply(apply, openedType, locator, ctorLocator);
+  auto *ctorCall = finishApply(apply, openedType, locator, ctorLocator);
+
+  // Check whether this is a situation like `T(...) { ... }` where `T` is
+  // a callable type and trailing closure(s) are associated with implicit
+  // `.callAsFunction` instead of constructor.
+  {
+    auto callAsFunction =
+        solution.ImplicitCallAsFunctionRoots.find(ctorLocator);
+    if (callAsFunction != solution.ImplicitCallAsFunctionRoots.end()) {
+      auto *dotExpr = callAsFunction->second;
+      auto resultTy = solution.getResolvedType(dotExpr);
+
+      auto *implicitCall = CallExpr::createImplicit(
+          cs.getASTContext(), ctorCall,
+          solution.getArgumentList(cs.getConstraintLocator(
+              dotExpr, ConstraintLocator::ApplyArgument)));
+
+      implicitCall->setType(resultTy);
+      cs.cacheType(implicitCall);
+
+      auto *memberCalleeLoc =
+          cs.getConstraintLocator(dotExpr,
+                                  {ConstraintLocator::ApplyFunction,
+                                   ConstraintLocator::ImplicitCallAsFunction},
+                                  /*summaryFlags=*/0);
+
+      return finishApply(implicitCall, resultTy, cs.getConstraintLocator(dotExpr),
+                         memberCalleeLoc);
+    }
+  }
+
+  return ctorCall;
 }
 
 /// Determine whether this closure should be treated as Sendable.

lib/Sema/CSSimplify.cpp

@@ -1321,6 +1321,7 @@ class OpenTypeSequenceElements {
 // Match the argument of a call to the parameter.
 ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
     ConstraintSystem &cs, FunctionType *contextualType,
+    ArgumentList *argList,
     ArrayRef<AnyFunctionType::Param> args,
     ArrayRef<AnyFunctionType::Param> params, ConstraintKind subKind,
     ConstraintLocatorBuilder locator,
@@ -1340,8 +1341,7 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
 
   ParameterListInfo paramInfo(params, callee, appliedSelf);
 
-  // Dig out the argument information.
-  auto *argList = cs.getArgumentList(loc);
+  // Make sure that argument list is available.
   assert(argList);
 
   // Apply labels to arguments.
@@ -10494,6 +10494,32 @@ bool ConstraintSystem::simplifyAppliedOverloads(
                                       numOptionalUnwraps, locator);
 }
 
+/// Create an implicit dot-member reference expression to be used
+/// as a root for injected `.callAsFunction` call.
+static UnresolvedDotExpr *
+createImplicitRootForCallAsFunction(ConstraintSystem &cs, Type refType,
+                                    ArgumentList *arguments,
+                                    ConstraintLocator *calleeLocator) {
+  auto &ctx = cs.getASTContext();
+  auto *baseExpr = castToExpr(calleeLocator->getAnchor());
+
+  SmallVector<Identifier, 2> closureLabelsScratch;
+  // Create implicit `.callAsFunction` expression to use as an anchor
+  // for new argument list that only has trailing closures in it.
+  auto *implicitRef = UnresolvedDotExpr::createImplicit(
+      ctx, baseExpr, {ctx.Id_callAsFunction},
+      arguments->getArgumentLabels(closureLabelsScratch));
+
+  {
+    // Record a type of the new reference in the constraint system.
+    cs.setType(implicitRef, refType);
+    // Record new `.callAsFunction` in the constraint system.
+    cs.recordCallAsFunction(implicitRef, arguments, calleeLocator);
+  }
+
+  return implicitRef;
+}
+
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyApplicableFnConstraint(
     Type type1, Type type2,
@@ -10548,17 +10574,20 @@ ConstraintSystem::simplifyApplicableFnConstraint(
   };
 
   // Local function to form an unsolved result.
-  auto formUnsolved = [&] {
+  auto formUnsolved = [&](bool activate = false) {
     if (flags.contains(TMF_GenerateConstraints)) {
-      addUnsolvedConstraint(
-        Constraint::createApplicableFunction(
+      auto *application = Constraint::createApplicableFunction(
           *this, type1, type2, trailingClosureMatching,
-          getConstraintLocator(locator)));
+          getConstraintLocator(locator));
+
+      addUnsolvedConstraint(application);
+      if (activate)
+        activateConstraint(application);
+
       return SolutionKind::Solved;
     }
 
     return SolutionKind::Unsolved;
-
   };
 
   // If right-hand side is a type variable, the constraint is unsolved.
@@ -10633,15 +10662,97 @@ ConstraintSystem::simplifyApplicableFnConstraint(
                               ? ConstraintKind::OperatorArgumentConversion
                               : ConstraintKind::ArgumentConversion);
 
+    auto *argumentsLoc = getConstraintLocator(
+        outerLocator.withPathElement(ConstraintLocator::ApplyArgument));
+
+    auto *argumentList = getArgumentList(argumentsLoc);
     // The argument type must be convertible to the input type.
     auto matchCallResult = ::matchCallArguments(
-        *this, func2, func1->getParams(), func2->getParams(), subKind,
-        outerLocator.withPathElement(ConstraintLocator::ApplyArgument),
-        trailingClosureMatching);
+        *this, func2, argumentList, func1->getParams(), func2->getParams(),
+        subKind, argumentsLoc, trailingClosureMatching);
 
     switch (matchCallResult) {
-    case SolutionKind::Error:
+    case SolutionKind::Error: {
+      auto resultTy = func2->getResult();
+
+      // If this is a call that constructs a callable type with
+      // trailing closure(s), closure(s) might not belong to
+      // the constructor but rather to implicit `callAsFunction`,
+      // there is no way to determine that without trying.
+      if (resultTy->isCallableNominalType(DC) &&
+          argumentList->hasAnyTrailingClosures()) {
+        auto *calleeLoc = getCalleeLocator(argumentsLoc);
+
+        bool isInit = false;
+        if (auto overload = findSelectedOverloadFor(calleeLoc)) {
+          isInit = bool(dyn_cast_or_null<ConstructorDecl>(
+              overload->choice.getDeclOrNull()));
+        }
+
+        if (!isInit)
+          return SolutionKind::Error;
+
+        auto &ctx = getASTContext();
+        auto numTrailing = argumentList->getNumTrailingClosures();
+
+        SmallVector<Argument, 4> newArguments(
+            argumentList->getNonTrailingArgs());
+        SmallVector<Argument, 4> trailingClosures(
+            argumentList->getTrailingClosures());
+
+        // Original argument list with all the trailing closures removed.
+        auto *newArgumentList = ArgumentList::createParsed(
+            ctx, argumentList->getLParenLoc(), newArguments,
+            argumentList->getRParenLoc(),
+            /*firstTrailingClosureIndex=*/None);
+
+        auto trailingClosureTypes = func1->getParams().take_back(numTrailing);
+        // The original result type is going to become a result of
+        // implicit `.callAsFunction` instead since `.callAsFunction`
+        // is inserted between `.init` and trailing closures.
+        auto callAsFunctionResultTy = func1->getResult();
+
+        // The implicit replacement for original result type which
+        // represents a callable type produced by `.init` call.
+        auto callableType =
+            createTypeVariable(getConstraintLocator({}), /*flags=*/0);
+
+        // The original application type with all the trailing closures
+        // dropped from it and result replaced to the implicit variable.
+        func1 = FunctionType::get(func1->getParams().drop_back(numTrailing),
+                                  callableType, func1->getExtInfo());
+
+        auto matchCallResult = ::matchCallArguments(
+            *this, func2, newArgumentList, func1->getParams(),
+            func2->getParams(), subKind, argumentsLoc, trailingClosureMatching);
+
+        if (matchCallResult != SolutionKind::Solved)
+          return SolutionKind::Error;
+
+        auto *implicitCallArgumentList =
+            ArgumentList::createImplicit(ctx, trailingClosures,
+                                         /*firstTrailingClosureIndex=*/0);
+
+        auto *implicitRef = createImplicitRootForCallAsFunction(
+            *this, callAsFunctionResultTy, implicitCallArgumentList, calleeLoc);
+
+        auto callAsFunctionArguments =
+            FunctionType::get(trailingClosureTypes, callAsFunctionResultTy,
+                              FunctionType::ExtInfo());
+
+        // Form an unsolved constraint to apply trailing closures to a
+        // callable type produced by `.init`. This constraint would become
+        // active when `callableType` is bound.
+        addUnsolvedConstraint(Constraint::create(
+            *this, ConstraintKind::ApplicableFunction, callAsFunctionArguments,
+            callableType,
+            getConstraintLocator(implicitRef,
+                                 ConstraintLocator::ApplyFunction)));
+        break;
+      }
+
       return SolutionKind::Error;
+    }
 
     case SolutionKind::Unsolved: {
       // Only occurs when there is an ambiguity between forward scanning and
@@ -10691,6 +10802,26 @@ ConstraintSystem::simplifyApplicableFnConstraint(
     if (instance2->isTypeVariableOrMember())
       return formUnsolved();
 
+    auto *argumentsLoc = getConstraintLocator(
+        outerLocator.withPathElement(ConstraintLocator::ApplyArgument));
+
+    auto *argumentList = getArgumentList(argumentsLoc);
+    assert(argumentList);
+
+    // Cannot simplify construction of callable types during constraint
+    // generation when trailing closures are present because such calls
+    // have special trailing closure matching semantics. It's unclear
+    // whether trailing arguments belong to `.init` or implicit
+    // `.callAsFunction` in this case.
+    //
+    // Note that the constraint has to be activate so that solver attempts
+    // once constraint generation is done.
+    if (getPhase() == ConstraintSystemPhase::ConstraintGeneration &&
+        argumentList->hasAnyTrailingClosures() &&
+        instance2->isCallableNominalType(DC)) {
+      return formUnsolved(/*activate=*/true);
+    }
+
     // Construct the instance from the input arguments.
     auto simplified = simplifyConstructionConstraint(instance2, func1, subflags,
                                           /*FIXME?*/ DC,
@@ -11579,6 +11710,7 @@ ConstraintSystem::simplifyRestrictedConstraintImpl(
     if (!ArgumentLists.count(memberLoc)) {
       auto *argList = ArgumentList::createImplicit(
           getASTContext(), {Argument(SourceLoc(), Identifier(), nullptr)},
+          /*firstTrailingClosureIndex=*/None,
           AllocationArena::ConstraintSolver);
       ArgumentLists.insert({memberLoc, argList});
     }
@@ -11867,6 +11999,15 @@ void ConstraintSystem::recordMatchCallArgumentResult(
   argumentMatchingChoices.insert({locator, result});
 }
 
+void ConstraintSystem::recordCallAsFunction(UnresolvedDotExpr *root,
+                                            ArgumentList *arguments,
+                                            ConstraintLocator *locator) {
+  ImplicitCallAsFunctionRoots.insert({locator, root});
+
+  associateArgumentList(
+      getConstraintLocator(root, ConstraintLocator::ApplyArgument), arguments);
+}
+
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
     ConstraintFix *fix, Type type1, Type type2, ConstraintKind matchKind,
     TypeMatchOptions flags, ConstraintLocatorBuilder locator) {