[CS] Store argument list mappings on solutions

Roll back argument list mappings in the constraint
system at the end of solver scopes, and copy
argument list mappings into solutions.

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -1198,6 +1198,11 @@ class Solution {
   /// A map from argument expressions to their applied property wrapper expressions.
   llvm::MapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>> appliedPropertyWrappers;
 
+  /// A mapping from the constraint locators for references to various
+  /// names (e.g., member references, normal name references, possible
+  /// constructions) to the argument lists for the call to that locator.
+  llvm::MapVector<ConstraintLocator *, ArgumentList *> argumentLists;
+
   /// Record a new argument matching choice for given locator that maps a
   /// single argument to a single parameter.
   void recordSingleArgMatchingChoice(ConstraintLocator *locator);
@@ -1332,6 +1337,10 @@ class Solution {
   /// expression type map.
   bool isStaticallyDerivedMetatype(Expr *E) const;
 
+  /// Retrieve the argument list that is associated with a call at the given
+  /// locator.
+  ArgumentList *getArgumentList(ConstraintLocator *locator) const;
+
   SWIFT_DEBUG_DUMP;
 
   /// Dump this solution.
@@ -2409,7 +2418,7 @@ class ConstraintSystem {
   /// A mapping from the constraint locators for references to various
   /// names (e.g., member references, normal name references, possible
   /// constructions) to the argument lists for the call to that locator.
-  llvm::DenseMap<ConstraintLocator *, ArgumentList *> ArgumentLists;
+  llvm::MapVector<ConstraintLocator *, ArgumentList *> ArgumentLists;
 
 public:
   /// A map from argument expressions to their applied property wrapper expressions.
@@ -2899,6 +2908,9 @@ class ConstraintSystem {
     /// The length of \c ImplicitValueConversions.
     unsigned numImplicitValueConversions;
 
+    /// The length of \c ArgumentLists.
+    unsigned numArgumentLists;
+
     /// The previous score.
     Score PreviousScore;
 

lib/Sema/CSDiagnostics.cpp

@@ -112,7 +112,7 @@ Expr *FailureDiagnostic::findParentExpr(const Expr *subExpr) const {
 
 ArgumentList *
 FailureDiagnostic::getArgumentListFor(ConstraintLocator *locator) const {
-  return getConstraintSystem().getArgumentList(locator);
+  return S.getArgumentList(locator);
 }
 
 Expr *FailureDiagnostic::getBaseExprFor(const Expr *anchor) const {

lib/Sema/CSSolver.cpp

@@ -199,6 +199,11 @@ Solution ConstraintSystem::finalize() {
     solution.ImplicitValueConversions.push_back(valueConversion);
   }
 
+  // Remember argument lists.
+  for (const auto &argListMapping : ArgumentLists) {
+    solution.argumentLists.insert(argListMapping);
+  }
+
   return solution;
 }
 
@@ -297,6 +302,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     ImplicitValueConversions.push_back(valueConversion);
   }
 
+  // Register the argument lists.
+  for (auto &argListMapping : solution.argumentLists) {
+    ArgumentLists.insert(argListMapping);
+  }
+
   // Register any fixes produced along this path.
   Fixes.append(solution.Fixes.begin(), solution.Fixes.end());
 }
@@ -510,6 +520,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numSolutionApplicationTargets = cs.solutionApplicationTargets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
+  numArgumentLists = cs.ArgumentLists.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -619,6 +630,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any implicit value conversions.
   truncate(cs.ImplicitValueConversions, numImplicitValueConversions);
 
+  // Remove any argument lists no longer in scope.
+  truncate(cs.ArgumentLists, numArgumentLists);
+
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;
 

lib/Sema/ConstraintSystem.cpp

@@ -3713,10 +3713,11 @@ static bool diagnoseAmbiguity(
                   /*isApplication=*/false, decl->getDescriptiveKind(),
                   name.isSpecial(), name.getBaseName());
     } else {
-      bool isApplication =
-          llvm::any_of(cs.ArgumentLists, [&](const auto &pair) {
+      bool isApplication = llvm::any_of(solutions, [&](const auto &S) {
+          return llvm::any_of(S.argumentLists, [&](const auto &pair) {
             return pair.first->getAnchor() == commonAnchor;
           });
+      });
 
       DE.diagnose(getLoc(commonAnchor),
                   diag::no_overloads_match_exactly_in_call, isApplication,
@@ -3757,7 +3758,7 @@ static bool diagnoseAmbiguity(
 
         if (fn->getNumParams() == 1) {
           auto *argList =
-              cs.getArgumentList(solution.Fixes.front()->getLocator());
+              solution.getArgumentList(solution.Fixes.front()->getLocator());
           assert(argList);
 
           const auto &param = fn->getParams()[0];
@@ -4671,6 +4672,18 @@ void ConstraintSystem::associateArgumentList(ConstraintLocator *locator,
   (void)inserted;
 }
 
+ArgumentList *Solution::getArgumentList(ConstraintLocator *locator) const {
+  if (!locator)
+    return nullptr;
+
+  if (auto *infoLocator = constraintSystem->getArgumentInfoLocator(locator)) {
+    auto known = argumentLists.find(infoLocator);
+    if (known != argumentLists.end())
+      return known->second;
+  }
+  return nullptr;
+}
+
 /// Given an apply expr, returns true if it is expected to have a direct callee
 /// overload, resolvable using `getChoiceFor`. Otherwise, returns false.
 static bool shouldHaveDirectCalleeOverload(const CallExpr *callExpr) {
@@ -4728,12 +4741,10 @@ Type Solution::resolveInterfaceType(Type type) const {
 
 Optional<FunctionArgApplyInfo>
 Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
-  auto &cs = getConstraintSystem();
-
   // It's only valid to use `&` in argument positions, but we need
   // to figure out exactly where it was used.
   if (auto *argExpr = getAsExpr<InOutExpr>(locator->getAnchor())) {
-    auto *argLoc = cs.getArgumentLocator(argExpr);
+    auto *argLoc = getConstraintSystem().getArgumentLocator(argExpr);
     assert(argLoc && "Incorrect use of `inout` expression");
     locator = argLoc;
   }
@@ -4766,7 +4777,7 @@ Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   if (!argExpr)
     return None;
 
-  auto *argList = cs.getArgumentList(argLocator);
+  auto *argList = getArgumentList(argLocator);
   if (!argList)
     return None;