[ConstraintSystem] Add storage for implicitly generated roots of .callAsFunction

Some implicit calls to `.callAsFunction` require that a new root
expression to be created for them in order to record argument list
and resolved overload choice.

Author: xedin

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;
 

lib/Sema/CSSolver.cpp

@@ -201,6 +201,10 @@ Solution ConstraintSystem::finalize() {
     solution.argumentLists.insert(argListMapping);
   }
 
+  for (const auto &implicitRoot : ImplicitCallAsFunctionRoots) {
+    solution.ImplicitCallAsFunctionRoots.insert(implicitRoot);
+  }
+
   return solution;
 }
 
@@ -308,6 +312,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     ArgumentLists.insert(argListMapping);
   }
 
+  for (auto &implicitRoot : solution.ImplicitCallAsFunctionRoots) {
+    ImplicitCallAsFunctionRoots.insert(implicitRoot);
+  }
+
   // Register any fixes produced along this path.
   Fixes.insert(solution.Fixes.begin(), solution.Fixes.end());
 }
@@ -529,6 +537,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numIsolatedParams = cs.isolatedParams.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
   numArgumentLists = cs.ArgumentLists.size();
+  numImplicitCallAsFunctionRoots = cs.ImplicitCallAsFunctionRoots.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -643,6 +652,10 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any argument lists no longer in scope.
   truncate(cs.ArgumentLists, numArgumentLists);
 
+  // Remove any implicitly generated root expressions for `.callAsFunction`
+  // which are no longer in scope.
+  truncate(cs.ImplicitCallAsFunctionRoots, numImplicitCallAsFunctionRoots);
+
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;
 

lib/Sema/ConstraintSystem.cpp

@@ -3239,7 +3239,8 @@ size_t Solution::getTotalMemory() const {
          ConstraintRestrictions.getMemorySize() +
          llvm::capacity_in_bytes(Fixes) + DisjunctionChoices.getMemorySize() +
          OpenedTypes.getMemorySize() + OpenedExistentialTypes.getMemorySize() +
-         (DefaultedConstraints.size() * sizeof(void *));
+         (DefaultedConstraints.size() * sizeof(void *)) +
+         ImplicitCallAsFunctionRoots.getMemorySize();
 }
 
 DeclContext *Solution::getDC() const { return constraintSystem->DC; }