[IDE] Set solution-specific variable types as interface types

Setting the interface type of a variable, just to reset it to a null type is actually really gross. But quite a few methods further down in the generation of code completion results (such as USR generation) need to get a variable’s type and passing them along in a separate map would be really invasive. So this seems like the least bad solution to me.

include/swift/IDE/TypeCheckCompletionCallback.h

@@ -82,6 +82,43 @@ void getSolutionSpecificVarTypes(
     const constraints::Solution &S,
     llvm::SmallDenseMap<const VarDecl *, Type> &Result);
 
+/// While this RAII is alive the interface types of the variables defined in
+/// \c SolutionSpecificVarTypes are temporarily set to the types in the map.
+/// Afterwards, their types are restored.
+struct WithSolutionSpecificVarTypesRAII {
+  llvm::SmallDenseMap<const VarDecl *, Type> RestoreVarTypes;
+
+  WithSolutionSpecificVarTypesRAII(
+      llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes) {
+    for (auto SolutionVarType : SolutionSpecificVarTypes) {
+      if (SolutionVarType.first->hasInterfaceType()) {
+        RestoreVarTypes[SolutionVarType.first] =
+            SolutionVarType.first->getInterfaceType();
+      } else {
+        RestoreVarTypes[SolutionVarType.first] = Type();
+      }
+      if (!SolutionVarType.second->hasArchetype()) {
+        setInterfaceType(const_cast<VarDecl *>(SolutionVarType.first),
+                         SolutionVarType.second);
+      } else {
+        setInterfaceType(const_cast<VarDecl *>(SolutionVarType.first),
+                         ErrorType::get(SolutionVarType.second));
+      }
+    }
+  }
+
+  ~WithSolutionSpecificVarTypesRAII() {
+    for (auto Var : RestoreVarTypes) {
+      setInterfaceType(const_cast<VarDecl *>(Var.first), Var.second);
+    }
+  }
+
+private:
+  /// Sets the interface type of \p VD, similar to \c VD->setInterfaceType
+  /// but also allows resetting the interface type of \p VD to null.
+  static void setInterfaceType(VarDecl *VD, Type Ty);
+};
+
 /// 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

@@ -351,6 +351,13 @@ void ArgumentTypeCheckCompletionCallback::deliverResults(
   }
 
   if (shouldPerformGlobalCompletion) {
+    llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+    if (!Results.empty()) {
+      SolutionSpecificVarTypes = Results[0].SolutionSpecificVarTypes;
+    }
+
+    WithSolutionSpecificVarTypesRAII VarTypes(SolutionSpecificVarTypes);
+
     for (auto &Result : Results) {
       ExpectedTypes.push_back(Result.ExpectedType);
       Lookup.setSolutionSpecificVarTypes(Result.SolutionSpecificVarTypes);

lib/IDE/ExprCompletion.cpp

@@ -105,6 +105,8 @@ void ExprTypeCheckCompletionCallback::deliverResults(
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
 
   for (auto &Result : Results) {
+    WithSolutionSpecificVarTypesRAII VarTypes(Result.SolutionSpecificVarTypes);
+
     Lookup.setExpectedTypes(ExpectedTypes,
                             Result.IsImplicitSingleExpressionReturn);
     Lookup.setCanCurrDeclContextHandleAsync(Result.IsInAsyncContext);

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -136,6 +136,11 @@ void swift::ide::getSolutionSpecificVarTypes(
   }
 }
 
+void WithSolutionSpecificVarTypesRAII::setInterfaceType(VarDecl *VD, Type Ty) {
+  VD->getASTContext().evaluator.cacheOutput(InterfaceTypeRequest{VD},
+                                            std::move(Ty));
+}
+
 bool swift::ide::isImplicitSingleExpressionReturn(ConstraintSystem &CS,
                                                   Expr *CompletionExpr) {
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);

test/IDE/complete_in_closures.swift

@@ -496,3 +496,16 @@ func testCompleteInMatchOfAssociatedValueInSwitchCase() {
 }
 
 }
+
+func testReferenceToVariableDefinedInClosure() {
+  func takeClosure(_ x: () -> Void) {}
+
+  takeClosure {
+    let item = "\(1)"
+    #^VARIABLE_DEFINED_IN_CLOSURE^#
+  }
+  // VARIABLE_DEFINED_IN_CLOSURE: Begin completions
+  // VARIABLE_DEFINED_IN_CLOSURE: Decl[LocalVar]/Local:               item[#String#]; name=item
+  // VARIABLE_DEFINED_IN_CLOSURE: End completions
+}
+

test/SourceKit/TypeContextInfo/typecontext_in_closure.swift

@@ -0,0 +1,14 @@
+// import SwiftUI
+
+struct ProgressView {
+  @ViewBuilder var body: Int {
+    // RUN: %sourcekitd-test -req=typecontextinfo -pos=%(line + 1):35 %s -- %s
+    grame(width: max(12345678901, 2))
+  }
+}
+
+func grame(width: Int?) -> Int {}
+
+@resultBuilder struct ViewBuilder {
+  public static func buildBlock(_ content: Int) -> Int
+}