[ConstraintSystem] Record parameter bindings in solutions (NFC)

This saves us from needing to re-match args to params in CSApply and is also
useful for a forthcoming change migrating code completion in argument position
to use the solver-based typeCheckForCodeCompletion api.

rdar://76581093

Author: nathawes

include/swift/Sema/ConstraintSystem.h

@@ -1111,6 +1111,43 @@ class SolutionApplicationTargetsKey {
   }
 };
 
+/// Describes the arguments to which a parameter binds.
+/// FIXME: This is an awful data structure. We want the equivalent of a
+/// TinyPtrVector for unsigned values.
+using ParamBinding = SmallVector<unsigned, 1>;
+
+/// The result of calling matchCallArguments().
+struct MatchCallArgumentResult {
+  /// The direction of trailing closure matching that was performed.
+  TrailingClosureMatching trailingClosureMatching;
+
+  /// The parameter bindings determined by the match.
+  SmallVector<ParamBinding, 4> parameterBindings;
+
+  /// When present, the forward and backward scans each produced a result,
+  /// and the parameter bindings are different. The primary result will be
+  /// forwarding, and this represents the backward binding.
+  Optional<SmallVector<ParamBinding, 4>> backwardParameterBindings;
+
+  friend bool operator==(const MatchCallArgumentResult &lhs,
+                         const MatchCallArgumentResult &rhs) {
+    if (lhs.trailingClosureMatching != rhs.trailingClosureMatching)
+      return false;
+    if (lhs.parameterBindings != rhs.parameterBindings)
+      return false;
+    return lhs.backwardParameterBindings == rhs.backwardParameterBindings;
+  }
+
+  /// Generate a result that maps the provided number of arguments to the same
+  /// number of parameters via forward match.
+  static MatchCallArgumentResult forArity(unsigned argCount) {
+    SmallVector<ParamBinding, 4> Bindings;
+    for (unsigned i : range(argCount))
+      Bindings.push_back({i});
+    return {TrailingClosureMatching::Forward, Bindings, None};
+  }
+};
+
 /// A complete solution to a constraint system.
 ///
 /// A solution to a constraint system consists of type variable bindings to
@@ -1159,9 +1196,9 @@ class Solution {
   llvm::SmallVector<ConstraintFix *, 4> Fixes;
 
   /// For locators associated with call expressions, the trailing closure
-  /// matching rule that was applied.
-  llvm::SmallMapVector<ConstraintLocator*, TrailingClosureMatching, 4>
-    trailingClosureMatchingChoices;
+  /// matching rule and parameter bindings that were applied.
+  llvm::SmallMapVector<ConstraintLocator *, MatchCallArgumentResult, 4>
+      argumentMatchingChoices;
 
   /// The set of disjunction choices used to arrive at this solution,
   /// which informs constraint application.
@@ -1203,6 +1240,10 @@ class Solution {
   /// A map from argument expressions to their applied property wrapper expressions.
   llvm::MapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>> appliedPropertyWrappers;
 
+  /// Record a new argument matching choice for given locator that maps a
+  /// single argument to a single parameter.
+  void recordSingleArgMatchingChoice(ConstraintLocator *locator);
+
   /// Simplify the given type by substituting all occurrences of
   /// type variables for their fixed types.
   Type simplifyType(Type type) const;
@@ -2210,9 +2251,9 @@ class ConstraintSystem {
       AppliedDisjunctions;
 
   /// For locators associated with call expressions, the trailing closure
-  /// matching rule that was applied.
-  std::vector<std::pair<ConstraintLocator*, TrailingClosureMatching>>
-      trailingClosureMatchingChoices;
+  /// matching rule and parameter bindings that were applied.
+  std::vector<std::pair<ConstraintLocator *, MatchCallArgumentResult>>
+      argumentMatchingChoices;
 
   /// The set of implicit value conversions performed by the solver on
   /// a current path to reach a solution.
@@ -2709,8 +2750,8 @@ class ConstraintSystem {
     /// The length of \c AppliedDisjunctions.
     unsigned numAppliedDisjunctions;
 
-    /// The length of \c trailingClosureMatchingChoices;
-    unsigned numTrailingClosureMatchingChoices;
+    /// The length of \c argumentMatchingChoices.
+    unsigned numArgumentMatchingChoices;
 
     /// The length of \c OpenedTypes.
     unsigned numOpenedTypes;
@@ -3226,11 +3267,8 @@ class ConstraintSystem {
 
   void recordPotentialHole(Type type);
 
-  void recordTrailingClosureMatch(
-      ConstraintLocator *locator,
-      TrailingClosureMatching trailingClosureMatch) {
-    trailingClosureMatchingChoices.push_back({locator, trailingClosureMatch});
-  }
+  void recordMatchCallArgumentResult(ConstraintLocator *locator,
+                                     MatchCallArgumentResult result);
 
   /// Walk a closure AST to determine its effects.
   ///
@@ -5086,11 +5124,6 @@ static inline bool computeTupleShuffle(TupleType *fromTuple,
                              sources);
 }
 
-/// Describes the arguments to which a parameter binds.
-/// FIXME: This is an awful data structure. We want the equivalent of a
-/// TinyPtrVector for unsigned values.
-using ParamBinding = SmallVector<unsigned, 1>;
-
 /// Class used as the base for listeners to the \c matchCallArguments process.
 ///
 /// By default, none of the callbacks do anything.
@@ -5158,20 +5191,6 @@ class MatchCallArgumentListener {
   virtual bool relabelArguments(ArrayRef<Identifier> newNames);
 };
 
-/// The result of calling matchCallArguments().
-struct MatchCallArgumentResult {
-  /// The direction of trailing closure matching that was performed.
-  TrailingClosureMatching trailingClosureMatching;
-
-  /// The parameter bindings determined by the match.
-  SmallVector<ParamBinding, 4> parameterBindings;
-
-  /// When present, the forward and backward scans each produced a result,
-  /// and the parameter bindings are different. The primary result will be
-  /// forwarding, and this represents the backward binding.
-  Optional<SmallVector<ParamBinding, 4>> backwardParameterBindings;
-};
-
 /// Match the call arguments (as described by the given argument type) to
 /// the parameters (as described by the given parameter type).
 ///

lib/Sema/CSApply.cpp

@@ -3099,16 +3099,16 @@ namespace {
                            DeclNameLoc nameLoc, bool implicit,
                            ConstraintLocator *ctorLocator,
                            SelectedOverload overload) {
+      auto locator = cs.getConstraintLocator(expr);
       auto choice = overload.choice;
       assert(choice.getKind() != OverloadChoiceKind::DeclViaDynamic);
       auto *ctor = cast<ConstructorDecl>(choice.getDecl());
 
       // If the subexpression is a metatype, build a direct reference to the
       // constructor.
       if (cs.getType(base)->is<AnyMetatypeType>()) {
-        return buildMemberRef(
-            base, dotLoc, overload, nameLoc, cs.getConstraintLocator(expr),
-            ctorLocator, implicit, AccessSemantics::Ordinary);
+        return buildMemberRef(base, dotLoc, overload, nameLoc, locator,
+                              ctorLocator, implicit, AccessSemantics::Ordinary);
       }
 
       // The subexpression must be either 'self' or 'super'.
@@ -3158,8 +3158,7 @@ namespace {
       auto *call = new (cs.getASTContext()) DotSyntaxCallExpr(ctorRef, dotLoc,
                                                               base);
 
-      return finishApply(call, cs.getType(expr), cs.getConstraintLocator(expr),
-                         ctorLocator);
+      return finishApply(call, cs.getType(expr), locator, ctorLocator);
     }
 
     /// Give the deprecation warning for referring to a global function
@@ -4841,19 +4840,19 @@ namespace {
         }
 
         auto kind = origComponent.getKind();
-        auto locator = cs.getConstraintLocator(
-            E, LocatorPathElt::KeyPathComponent(i));
+        auto componentLocator =
+            cs.getConstraintLocator(E, LocatorPathElt::KeyPathComponent(i));
 
-        // Adjust the locator such that it includes any additional elements to
-        // point to the component's callee, e.g a SubscriptMember for a
-        // subscript component.
-        locator = cs.getCalleeLocator(locator);
+        // Get a locator such that it includes any additional elements to point
+        // to the component's callee, e.g a SubscriptMember for a subscript
+        // component.
+        auto calleeLoc = cs.getCalleeLocator(componentLocator);
 
         bool isDynamicMember = false;
         // If this is an unresolved link, make sure we resolved it.
         if (kind == KeyPathExpr::Component::Kind::UnresolvedProperty ||
             kind == KeyPathExpr::Component::Kind::UnresolvedSubscript) {
-          auto foundDecl = solution.getOverloadChoiceIfAvailable(locator);
+          auto foundDecl = solution.getOverloadChoiceIfAvailable(calleeLoc);
           if (!foundDecl) {
             // If we couldn't resolve the component, leave it alone.
             resolvedComponents.push_back(origComponent);
@@ -4876,9 +4875,9 @@ namespace {
 
         switch (kind) {
         case KeyPathExpr::Component::Kind::UnresolvedProperty: {
-          buildKeyPathPropertyComponent(solution.getOverloadChoice(locator),
-                                        origComponent.getLoc(),
-                                        locator, resolvedComponents);
+          buildKeyPathPropertyComponent(solution.getOverloadChoice(calleeLoc),
+                                        origComponent.getLoc(), calleeLoc,
+                                        resolvedComponents);
           break;
         }
         case KeyPathExpr::Component::Kind::UnresolvedSubscript: {
@@ -4887,9 +4886,9 @@ namespace {
             subscriptLabels = origComponent.getSubscriptLabels();
 
           buildKeyPathSubscriptComponent(
-              solution.getOverloadChoice(locator),
-              origComponent.getLoc(), origComponent.getIndexExpr(),
-              subscriptLabels, locator, resolvedComponents);
+              solution.getOverloadChoice(calleeLoc), origComponent.getLoc(),
+              origComponent.getIndexExpr(), subscriptLabels, componentLocator,
+              resolvedComponents);
           break;
         }
         case KeyPathExpr::Component::Kind::OptionalChain: {
@@ -5138,9 +5137,10 @@ namespace {
         SmallVectorImpl<KeyPathExpr::Component> &components) {
       auto subscript = cast<SubscriptDecl>(overload.choice.getDecl());
       assert(!subscript->isGetterMutating());
+      auto memberLoc = cs.getCalleeLocator(locator);
 
       // Compute substitutions to refer to the member.
-      auto ref = resolveConcreteDeclRef(subscript, locator);
+      auto ref = resolveConcreteDeclRef(subscript, memberLoc);
 
       // If this is a @dynamicMemberLookup reference to resolve a property
       // through the subscript(dynamicMember:) member, restore the
@@ -5159,23 +5159,27 @@ namespace {
         if (overload.choice.getKind() ==
             OverloadChoiceKind::KeyPathDynamicMemberLookup) {
           indexExpr = buildKeyPathDynamicMemberIndexExpr(
-              indexType->castTo<BoundGenericType>(), componentLoc, locator);
+              indexType->castTo<BoundGenericType>(), componentLoc, memberLoc);
         } else {
           auto fieldName = overload.choice.getName().getBaseIdentifier().str();
           indexExpr = buildDynamicMemberLookupIndexExpr(fieldName, componentLoc,
                                                         indexType);
         }
+        // Record the implicit subscript expr's parameter bindings and matching
+        // direction as `coerceCallArguments` requires them.
+        solution.recordSingleArgMatchingChoice(locator);
       }
 
       auto subscriptType =
           simplifyType(overload.openedType)->castTo<AnyFunctionType>();
       auto resolvedTy = subscriptType->getResult();
 
       // Coerce the indices to the type the subscript expects.
-      auto *newIndexExpr =
-          coerceCallArguments(indexExpr, subscriptType, ref,
-                              /*applyExpr*/ nullptr, labels,
-                              locator, /*appliedPropertyWrappers*/ {});
+      auto *newIndexExpr = coerceCallArguments(
+          indexExpr, subscriptType, ref,
+          /*applyExpr*/ nullptr, labels,
+          cs.getConstraintLocator(locator, ConstraintLocator::ApplyArgument),
+          /*appliedPropertyWrappers*/ {});
 
       // We need to be able to hash the captured index values in order for
       // KeyPath itself to be hashable, so check that all of the subscript
@@ -5761,6 +5765,8 @@ Expr *ExprRewriter::coerceCallArguments(
     ArrayRef<Identifier> argLabels,
     ConstraintLocatorBuilder locator,
     ArrayRef<AppliedPropertyWrapper> appliedPropertyWrappers) {
+  assert(locator.last() && locator.last()->is<LocatorPathElt::ApplyArgument>());
+
   auto &ctx = getConstraintSystem().getASTContext();
   auto params = funcType->getParams();
   unsigned appliedWrapperIndex = 0;
@@ -5773,11 +5779,6 @@ Expr *ExprRewriter::coerceCallArguments(
         LocatorPathElt::ApplyArgToParam(argIdx, paramIdx, flags));
   };
 
-  bool matchCanFail =
-      llvm::any_of(params, [](const AnyFunctionType::Param &param) {
-        return param.getPlainType()->hasUnresolvedType();
-      });
-
   // Determine whether this application has curried self.
   bool skipCurriedSelf = apply ? hasCurriedSelf(cs, callee, apply) : true;
   // Determine the parameter bindings.
@@ -5812,34 +5813,14 @@ Expr *ExprRewriter::coerceCallArguments(
   // Apply labels to arguments.
   AnyFunctionType::relabelParams(args, argLabels);
 
-  MatchCallArgumentListener listener;
   auto unlabeledTrailingClosureIndex =
       arg->getUnlabeledTrailingClosureIndexOfPackedArgument();
 
-  // Determine the trailing closure matching rule that was applied. This
-  // is only relevant for explicit calls and subscripts.
-  auto trailingClosureMatching = TrailingClosureMatching::Forward;
-  {
-    SmallVector<LocatorPathElt, 4> path;
-    auto anchor = locator.getLocatorParts(path);
-    if (!path.empty() && path.back().is<LocatorPathElt::ApplyArgument>() &&
-        !anchor.isExpr(ExprKind::UnresolvedDot)) {
-      auto locatorPtr = cs.getConstraintLocator(locator);
-      assert(solution.trailingClosureMatchingChoices.count(locatorPtr) == 1);
-      trailingClosureMatching = solution.trailingClosureMatchingChoices.find(
-          locatorPtr)->second;
-    }
-  }
-
-  auto callArgumentMatch = constraints::matchCallArguments(
-      args, params, paramInfo, unlabeledTrailingClosureIndex,
-      /*allowFixes=*/false, listener, trailingClosureMatching);
-
-  assert((matchCanFail || callArgumentMatch) &&
-         "Call arguments did not match up?");
-  (void)matchCanFail;
-
-  auto parameterBindings = std::move(callArgumentMatch->parameterBindings);
+  // Determine the parameter bindings that were applied.
+  auto *locatorPtr = cs.getConstraintLocator(locator);
+  assert(solution.argumentMatchingChoices.count(locatorPtr) == 1);
+  auto parameterBindings = solution.argumentMatchingChoices.find(locatorPtr)
+                               ->second.parameterBindings;
 
   // We should either have parentheses or a tuple.
   auto *argTuple = dyn_cast<TupleExpr>(arg);
@@ -6849,12 +6830,11 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
                         ConstraintLocator::ConstructorMember}));
 
         solution.overloadChoices.insert({memberLoc, overload});
-        solution.trailingClosureMatchingChoices.insert(
-            {cs.getConstraintLocator(callLocator,
-                                     ConstraintLocator::ApplyArgument),
-             TrailingClosureMatching::Forward});
       }
 
+      // Record the implicit call's parameter bindings and match direction.
+      solution.recordSingleArgMatchingChoice(callLocator);
+
       finishApply(implicitInit, toType, callLocator, callLocator);
       return implicitInit;
     }

lib/Sema/CSGen.cpp

@@ -1477,6 +1477,14 @@ namespace {
     }
 
     Type visitUnresolvedDotExpr(UnresolvedDotExpr *expr) {
+      // UnresolvedDot applies the base to remove a single curry level from a
+      // member reference without using an applicable function constraint so
+      // we record the call argument matching here so it can be found later when
+      // a solution is applied to the AST.
+      CS.recordMatchCallArgumentResult(
+          CS.getConstraintLocator(expr, ConstraintLocator::ApplyArgument),
+          MatchCallArgumentResult::forArity(1));
+
       // If this is Builtin.type_join*, just return any type and move
       // on since we're going to discard this, and creating any type
       // variables for the reference will cause problems.

lib/Sema/CSSimplify.cpp

@@ -1309,9 +1309,9 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
     }
 
     selectedTrailingMatching = callArgumentMatch->trailingClosureMatching;
-    // Record the direction of matching used for this call.
-    cs.recordTrailingClosureMatch(cs.getConstraintLocator(locator),
-                                  selectedTrailingMatching);
+    // Record the matching direction and parameter bindings used for this call.
+    cs.recordMatchCallArgumentResult(cs.getConstraintLocator(locator),
+                                     *callArgumentMatch);
 
     // If there was a disjunction because both forward and backward were
     // possible, increase the score for forward matches to bias toward the
@@ -9716,10 +9716,10 @@ ConstraintSystem::simplifyApplicableFnConstraint(
   // have an explicit inout argument.
   if (type1.getPointer() == desugar2) {
     if (!isOperator || !hasInOut()) {
-      recordTrailingClosureMatch(
+      recordMatchCallArgumentResult(
           getConstraintLocator(
               outerLocator.withPathElement(ConstraintLocator::ApplyArgument)),
-          TrailingClosureMatching::Forward);
+          MatchCallArgumentResult::forArity(func1->getNumParams()));
       return SolutionKind::Solved;
     }
   }
@@ -10864,6 +10864,12 @@ void ConstraintSystem::recordPotentialHole(Type type) {
   });
 }
 
+void ConstraintSystem::recordMatchCallArgumentResult(
+    ConstraintLocator *locator, MatchCallArgumentResult result) {
+  assert(locator->isLastElement<LocatorPathElt::ApplyArgument>());
+  argumentMatchingChoices.push_back({locator, result});
+}
+
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
     ConstraintFix *fix, Type type1, Type type2, ConstraintKind matchKind,
     TypeMatchOptions flags, ConstraintLocatorBuilder locator) {