Merge pull request swiftlang#39192 from hamishknight/a-range-of-arguments

[CS] Store argument list mappings on 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.
@@ -2406,6 +2415,11 @@ class ConstraintSystem {
   /// Cache of the effects any closures visited.
   llvm::SmallDenseMap<ClosureExpr *, FunctionType::ExtInfo, 4> closureEffectsCache;
 
+  /// 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;
+
 public:
   /// A map from argument expressions to their applied property wrapper expressions.
   llvm::SmallMapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>, 4> appliedPropertyWrappers;
@@ -2784,20 +2798,18 @@ class ConstraintSystem {
   /// we're exploring.
   SolverState *solverState = nullptr;
 
-  /// 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;
-
   /// Form a locator that can be used to retrieve argument information cached in
   /// the constraint system for the callee described by the anchor of the
   /// passed locator.
   ConstraintLocator *getArgumentInfoLocator(ConstraintLocator *locator);
 
-  /// Retrieve the argument list that is associated with a member
-  /// reference at the given locator.
+  /// Retrieve the argument list that is associated with a call at the given
+  /// locator.
   ArgumentList *getArgumentList(ConstraintLocator *locator);
 
+  /// Associate an argument list with a call at a given locator.
+  void associateArgumentList(ConstraintLocator *locator, ArgumentList *args);
+
   Optional<SelectedOverload>
   findSelectedOverloadFor(ConstraintLocator *locator) const {
     auto result = ResolvedOverloads.find(locator);
@@ -2896,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/AST/ArgumentList.cpp

@@ -116,7 +116,7 @@ ArgumentList *ArgumentList::createImplicit(ASTContext &ctx, SourceLoc lParenLoc,
 ArgumentList *ArgumentList::createImplicit(ASTContext &ctx,
                                            ArrayRef<Argument> args,
                                            AllocationArena arena) {
-  return createImplicit(ctx, SourceLoc(), args, SourceLoc());
+  return createImplicit(ctx, SourceLoc(), args, SourceLoc(), arena);
 }
 
 ArgumentList *ArgumentList::forImplicitSingle(ASTContext &ctx, Identifier label,

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/CSGen.cpp

@@ -917,7 +917,7 @@ namespace {
         CS.getConstraintLocator(locator,
                                 ConstraintLocator::FunctionResult);
 
-      associateArgumentList(memberLocator, argList);
+      CS.associateArgumentList(memberLocator, argList);
 
       Type outputTy;
 
@@ -1174,7 +1174,7 @@ namespace {
 
     Type visitObjectLiteralExpr(ObjectLiteralExpr *expr) {
       auto *exprLoc = CS.getConstraintLocator(expr);
-      associateArgumentList(exprLoc, expr->getArgs());
+      CS.associateArgumentList(exprLoc, expr->getArgs());
 
       // If the expression has already been assigned a type; just use that type.
       if (expr->getType())
@@ -2687,7 +2687,7 @@ namespace {
     Type visitApplyExpr(ApplyExpr *expr) {
       auto fnExpr = expr->getFn();
 
-      associateArgumentList(CS.getConstraintLocator(expr), expr->getArgs());
+      CS.associateArgumentList(CS.getConstraintLocator(expr), expr->getArgs());
 
       if (auto *UDE = dyn_cast<UnresolvedDotExpr>(fnExpr)) {
         auto typeOperation = getTypeOperation(UDE, CS.getASTContext());
@@ -3442,11 +3442,6 @@ namespace {
       }
       llvm_unreachable("unhandled operation");
     }
-
-    void associateArgumentList(ConstraintLocator *locator, ArgumentList *args) {
-      assert(locator && locator->getAnchor());
-      CS.ArgumentLists[CS.getArgumentInfoLocator(locator)] = args;
-    }
   };
 
   class ConstraintWalker : public ASTWalker {

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];
@@ -4668,6 +4669,27 @@ ArgumentList *ConstraintSystem::getArgumentList(ConstraintLocator *locator) {
   return nullptr;
 }
 
+void ConstraintSystem::associateArgumentList(ConstraintLocator *locator,
+                                             ArgumentList *args) {
+  assert(locator && locator->getAnchor());
+  auto *argInfoLoc = getArgumentInfoLocator(locator);
+  auto inserted = ArgumentLists.insert({argInfoLoc, args}).second;
+  assert(inserted && "Multiple argument lists at 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) {
@@ -4725,12 +4747,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;
   }
@@ -4763,7 +4783,7 @@ Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   if (!argExpr)
     return None;
 
-  auto *argList = cs.getArgumentList(argLocator);
+  auto *argList = getArgumentList(argLocator);
   if (!argList)
     return None;