[CodeCompletion] Record solution-specific variable types in argument completion

We are doing global completion for function arguments, so we should also record the variable types determined by the solution.

Author: ahoppen

include/swift/IDE/ArgumentCompletion.h

@@ -48,6 +48,11 @@ class ArgumentTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
     /// Whether the surrounding context is async and thus calling async
     /// functions is supported.
     bool IsInAsyncContext;
+
+    /// Types of variables that were determined in the solution that produced
+    /// this result. This in particular includes parameters of closures that
+    /// were type-checked with the code completion expression.
+    llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
   };
 
   CodeCompletionExpr *CompletionExpr;

include/swift/IDE/TypeCheckCompletionCallback.h

@@ -75,6 +75,13 @@ Type getTypeForCompletion(const constraints::Solution &S, ASTNode Node);
 /// completion expression \p E.
 Type getPatternMatchType(const constraints::Solution &S, Expr *E);
 
+/// Populate \p Result with types of variables that were determined in the
+/// solution \p S. This in particular includes parameters of closures that
+/// were type-checked with the code completion expression.
+void getSolutionSpecificVarTypes(
+    const constraints::Solution &S,
+    llvm::SmallDenseMap<const VarDecl *, Type> &Result);
+
 /// Whether the given completion expression is the only expression in its
 /// containing closure or function body and its value is implicitly returned.
 ///

lib/IDE/ArgumentCompletion.cpp

@@ -193,9 +193,13 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
     return;
   }
 
+  llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+  getSolutionSpecificVarTypes(S, SolutionSpecificVarTypes);
+
   Results.push_back({ExpectedTy, isa<SubscriptExpr>(ParentCall), FuncD, FuncTy,
                      ArgIdx, ParamIdx, std::move(ClaimedParams),
-                     IsNoninitialVariadic, CallBaseTy, HasLabel, IsAsync});
+                     IsNoninitialVariadic, CallBaseTy, HasLabel, IsAsync,
+                     SolutionSpecificVarTypes});
 }
 
 void ArgumentTypeCheckCompletionCallback::deliverResults(
@@ -267,6 +271,7 @@ void ArgumentTypeCheckCompletionCallback::deliverResults(
   if (shouldPerformGlobalCompletion) {
     for (auto &Result : Results) {
       ExpectedTypes.push_back(Result.ExpectedType);
+      Lookup.setSolutionSpecificVarTypes(Result.SolutionSpecificVarTypes);
     }
     Lookup.setExpectedTypes(ExpectedTypes, false);
     bool IsInAsyncContext = llvm::any_of(

lib/IDE/ExprCompletion.cpp

@@ -78,11 +78,7 @@ void ExprTypeCheckCompletionCallback::sawSolutionImpl(
   bool IsAsync = isContextAsync(S, DC);
 
   llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
-  for (auto NT : S.nodeTypes) {
-    if (auto VD = dyn_cast_or_null<VarDecl>(NT.first.dyn_cast<Decl *>())) {
-      SolutionSpecificVarTypes[VD] = S.simplifyType(NT.second);
-    }
-  }
+  getSolutionSpecificVarTypes(S, SolutionSpecificVarTypes);
 
   addResult(ImplicitReturn, IsAsync, SolutionSpecificVarTypes);
   addExpectedType(ExpectedTy);

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -125,6 +125,17 @@ Type swift::ide::getPatternMatchType(const constraints::Solution &S, Expr *E) {
   return nullptr;
 }
 
+void swift::ide::getSolutionSpecificVarTypes(
+    const constraints::Solution &S,
+    llvm::SmallDenseMap<const VarDecl *, Type> &Result) {
+  assert(Result.empty());
+  for (auto NT : S.nodeTypes) {
+    if (auto VD = dyn_cast_or_null<VarDecl>(NT.first.dyn_cast<Decl *>())) {
+      Result[VD] = S.simplifyType(NT.second);
+    }
+  }
+}
+
 bool swift::ide::isImplicitSingleExpressionReturn(ConstraintSystem &CS,
                                                   Expr *CompletionExpr) {
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);