[CursorInfo] Fix a bug that caused incorrect discriminators to be assigned to variables inside closures

Author: ahoppen

lib/IDE/CursorInfo.cpp

@@ -37,11 +37,20 @@ void typeCheckDeclAndParentClosures(ValueDecl *VD) {
   // encoded in the USR of ParentContexts in the cursor info response.
   auto DC = VD->getDeclContext();
   while (DC->getParent()) {
-    if (auto Closure = dyn_cast<AbstractClosureExpr>(DC)) {
-      if (Closure->getType().isNull()) {
+    if (auto AbstractClosure = dyn_cast<AbstractClosureExpr>(DC)) {
+      if (AbstractClosure->getType().isNull()) {
         typeCheckASTNodeAtLoc(
             TypeCheckASTNodeAtLocContext::declContext(DC->getParent()),
-            Closure->getLoc());
+            AbstractClosure->getLoc());
+      }
+      if (auto Closure = dyn_cast<ClosureExpr>(AbstractClosure)) {
+        // Set the closure's body state to `TypeCheckedWithSignature`. This
+        // matches what the type checker is doing when applying the solution and
+        // is necessary so the USRs we compute in solver-based cursor info match
+        // those in the AST - the critical point is that
+        // `isSeparatelyTypeChecked` in `SetLocalDiscriminators::walkToExprPre`
+        // needs to evaluate to `false`.
+        Closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
       }
     }
     DC = DC->getParent();

test/SourceKit/CursorInfo/variable_discriminator_in_closure.swift

@@ -0,0 +1,19 @@
+func sink(_ a: (Int) -> Void) {}
+    
+func testSingleStatementClosure() {
+  sink { items in
+    // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):9 %s -- %s | %FileCheck %s --check-prefix=SINGLE-STMT-CLOSURE
+    var items = items
+  }
+  // SINGLE-STMT-CLOSURE: s:33variable_discriminator_in_closure26testSingleStatementClosureyyFySiXEfU_5itemsL_Sivp
+}
+
+
+func testMultiStatementClosure() {
+  sink { items in
+    // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):9 %s -- %s | %FileCheck %s --check-prefix=MULTI-STMT-CLOSURE
+    var items = items
+    print("xxx")
+  }
+  // MULTI-STMT-CLOSURE: s:33variable_discriminator_in_closure25testMultiStatementClosureyyFySiXEfU_5itemsL_Sivp
+}