[SoruceKit] Add the underlying variable to related identifiers for a captured variable

Previously, the related idents request wouldn’t look through caputred variables like `[foo]`. Change the logic to consider the captured variable as well as the variable that’s implicitly declared for use inside the closure.

rdar://81628899

Author: ahoppen

test/SourceKit/RelatedIdents/related_idents.swift

@@ -75,6 +75,24 @@ func `escapedName`(`x`: Int) {}
 escapedName(`x`: 2)
 escapedName(`x`:)
 
+func closureCapture() {
+  let test = 0
+
+  let outerClosure = { [test] in
+    let innerClosure = { [test] in
+      print(test)
+    }
+  }
+}
+
+func closureCaptureWithExplicitName() {
+  let test = 0
+
+  let closure = { [test = test] in
+    print(test)
+  }
+}
+
 // RUN: %sourcekitd-test -req=related-idents -pos=6:17 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK1 %s
 // CHECK1: START RANGES
 // CHECK1-NEXT: 1:7 - 2
@@ -165,3 +183,28 @@ escapedName(`x`:)
 // CHECK11-NEXT: 75:1 - 11
 // CHECK11-NEXT: 76:1 - 11
 
+
+// RUN: %sourcekitd-test -req=related-idents -pos=79:7 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK12 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=81:25 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK12 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=82:27 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK12 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=83:13 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK12 %s
+// CHECK12:      START RANGES
+// CHECK12-NEXT: 79:7 - 4
+// CHECK12-NEXT: 81:25 - 4
+// CHECK12-NEXT: 82:27 - 4
+// CHECK12-NEXT: 83:13 - 4
+// CHECK12-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=89:7 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK13 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=91:27 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK13 %s
+// CHECK13:      START RANGES
+// CHECK13-NEXT: 89:7 - 4
+// CHECK13-NEXT: 91:27 - 4
+// CHECK13-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=91:20 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK14 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=92:11 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK14 %s
+// CHECK14:      START RANGES
+// CHECK14-NEXT: 91:20 - 4
+// CHECK14-NEXT: 92:11 - 4
+// CHECK14-NEXT: END RANGES

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2158,17 +2158,22 @@ SwiftLangSupport::findUSRRange(StringRef DocumentName, StringRef USR) {
 namespace {
 class RelatedIdScanner : public SourceEntityWalker {
   ValueDecl *Dcl;
-  llvm::SmallVectorImpl<std::pair<unsigned, unsigned>> &Ranges;
+  llvm::SmallDenseSet<std::pair<unsigned, unsigned>, 8> &Ranges;
+  /// Declarations that are tied to the same name as \c Dcl and should thus also
+  /// be renamed if \c Dcl is renamed. Most notabliy this contains closure
+  /// captures like `[foo]`.
+  llvm::SmallVectorImpl<ValueDecl *> &RelatedDecls;
   SourceManager &SourceMgr;
   unsigned BufferID = -1;
   bool Cancelled = false;
 
 public:
-  explicit RelatedIdScanner(SourceFile &SrcFile, unsigned BufferID,
-                            ValueDecl *D,
-      llvm::SmallVectorImpl<std::pair<unsigned, unsigned>> &Ranges)
-    : Ranges(Ranges), SourceMgr(SrcFile.getASTContext().SourceMgr),
-      BufferID(BufferID) {
+  explicit RelatedIdScanner(
+      SourceFile &SrcFile, unsigned BufferID, ValueDecl *D,
+      llvm::SmallDenseSet<std::pair<unsigned, unsigned>, 8> &Ranges,
+      llvm::SmallVectorImpl<ValueDecl *> &RelatedDecls)
+      : Ranges(Ranges), RelatedDecls(RelatedDecls),
+        SourceMgr(SrcFile.getASTContext().SourceMgr), BufferID(BufferID) {
     if (auto *V = dyn_cast<VarDecl>(D)) {
       // Always use the canonical var decl for comparison. This is so we
       // pick up all occurrences of x in case statements like the below:
@@ -2190,6 +2195,45 @@ class RelatedIdScanner : public SourceEntityWalker {
   }
 
 private:
+  bool walkToExprPre(Expr *E) override {
+    if (Cancelled)
+      return false;
+
+    // Check if there are closure captures like `[foo]` where the caputred
+    // variable should also be renamed
+    if (auto CaptureList = dyn_cast<CaptureListExpr>(E)) {
+      for (auto Capture : CaptureList->getCaptureList()) {
+        if (Capture.PBD->getPatternList().size() != 1) {
+          continue;
+        }
+        auto *DRE = dyn_cast_or_null<DeclRefExpr>(Capture.PBD->getInit(0));
+        if (!DRE) {
+          continue;
+        }
+
+        auto DeclaredVar = Capture.getVar();
+        if (DeclaredVar->getLoc() != DRE->getLoc()) {
+          // We have a capture like `[foo]` if the declared var and the
+          // reference share the same location.
+          continue;
+        }
+
+        auto *ReferencedVar = dyn_cast_or_null<VarDecl>(DRE->getDecl());
+        if (!ReferencedVar) {
+          continue;
+        }
+
+        assert(DeclaredVar->getName() == ReferencedVar->getName());
+        if (DeclaredVar == Dcl) {
+          RelatedDecls.push_back(ReferencedVar);
+        } else if (ReferencedVar == Dcl) {
+          RelatedDecls.push_back(DeclaredVar);
+        }
+      }
+    }
+    return true;
+  }
+
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
     if (Cancelled)
       return false;
@@ -2232,7 +2276,7 @@ class RelatedIdScanner : public SourceEntityWalker {
 
   bool passId(CharSourceRange Range) {
     unsigned Offset = SourceMgr.getLocOffsetInBuffer(Range.getStart(),BufferID);
-    Ranges.push_back({ Offset, Range.getByteLength() });
+    Ranges.insert({Offset, Range.getByteLength()});
     return !Cancelled;
   }
 };
@@ -2301,21 +2345,54 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
         if (VD->isOperator())
           return;
 
-        RelatedIdScanner Scanner(SrcFile, BufferID, VD, Ranges);
+        // Record ranges in a set first so we don't record some ranges twice.
+        // This could happen in capture lists where e.g. `[foo]` is both the
+        // reference of the captured variable and the declaration of the
+        // variable usable in the closure.
+        llvm::SmallDenseSet<std::pair<unsigned, unsigned>, 8> RangesSet;
+
+        // List of decls whose ranges should be reported as related identifiers.
+        SmallVector<ValueDecl *, 2> Worklist;
+        Worklist.push_back(VD);
+
+        // Decls that we have already visited, so we don't walk circles.
+        SmallPtrSet<ValueDecl *, 2> VisitedDecls;
+        while (!Worklist.empty()) {
+          ValueDecl *Dcl = Worklist.back();
+          Worklist.pop_back();
+          if (!VisitedDecls.insert(Dcl).second) {
+            // We have already visited this decl. Don't visit it again.
+            continue;
+          }
 
-        if (auto *Case = getCaseStmtOfCanonicalVar(VD)) {
-          Scanner.walk(Case);
-          while ((Case = Case->getFallthroughDest().getPtrOrNull())) {
+          RelatedIdScanner Scanner(SrcFile, BufferID, Dcl, RangesSet, Worklist);
+
+          if (auto *Case = getCaseStmtOfCanonicalVar(Dcl)) {
             Scanner.walk(Case);
+            while ((Case = Case->getFallthroughDest().getPtrOrNull())) {
+              Scanner.walk(Case);
+            }
+          } else if (DeclContext *LocalDC =
+                         Dcl->getDeclContext()->getLocalContext()) {
+            Scanner.walk(LocalDC);
+          } else {
+            Scanner.walk(SrcFile);
           }
-        } else if (DeclContext *LocalDC = VD->getDeclContext()->getLocalContext()) {
-          Scanner.walk(LocalDC);
-        } else {
-          Scanner.walk(SrcFile);
         }
+
+        // Sort ranges so we get deterministic output.
+        Ranges.insert(Ranges.end(), RangesSet.begin(), RangesSet.end());
+        llvm::sort(Ranges,
+                   [](const std::pair<unsigned, unsigned> &LHS,
+                      const std::pair<unsigned, unsigned> &RHS) -> bool {
+                     if (LHS.first == RHS.first) {
+                       return LHS.second < RHS.second;
+                     } else {
+                       return LHS.first < RHS.first;
+                     }
+                   });
       };
       Action();
-
       RelatedIdentsInfo Info;
       Info.Ranges = Ranges;
       Receiver(RequestResult<RelatedIdentsInfo>::fromResult(Info));