[sourcekitd][RelatedIdents] Fix handling of var decls in fallthrough case statements

We weren't picking up all occurrences of 'x' in the cases like the below:

case .first(let x), .second(let x):
  print("foo \(x)")
  fallthrough
case .third(let x):
  print("bar \(x)")

We would previously only return occurrences within the case statement the query
was made in (ignoring fallthroughs) and for cases with multiple patterns (as in
the first case above) we would only return the occurrence in the first pattern.

Author: Nathan Hawes

test/SourceKit/RelatedIdents/related_idents.swift

@@ -23,6 +23,27 @@ class C2<Param> {
 	func f(t : Param) -> Param {return t}
 }
 
+enum X {
+  case first(Int, String)
+  case second(Int, String)
+  case third(Int, String)
+  case fourth(Int, String)
+}
+
+switch X.first(2, "") {
+  case .first(let x, let y):
+    print(y)
+    fallthrough
+  case .second(let x, _):
+    print(x)
+  case .third(let x, let y):
+    fallthrough
+  case .fourth(let x, let y):
+    print(y)
+    print(x)
+    break
+}
+
 // 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
@@ -53,3 +74,37 @@ class C2<Param> {
 // CHECK5-NEXT: 23:13 - 5
 // CHECK5-NEXT: 23:23 - 5
 // CHECK5-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=34:19 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK6 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=37:20 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK6 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=38:11 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK6 %s
+// CHECK6:      START RANGES
+// CHECK6-NEXT: 34:19 - 1
+// CHECK6-NEXT: 37:20 - 1
+// CHECK6-NEXT: 38:11 - 1
+// CHECK6-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=34:26 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK7 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=35:11 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK7 %s
+// CHECK7:      START RANGES
+// CHECK7-NEXT: 34:26 - 1
+// CHECK7-NEXT: 35:11 - 1
+// CHECK7-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=39:26 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK8 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=41:27 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK8 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=42:11 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK8 %s
+// CHECK8:      START RANGES
+// CHECK8-NEXT: 39:26 - 1
+// CHECK8-NEXT: 41:27 - 1
+// CHECK8-NEXT: 42:11 - 1
+// CHECK8-NEXT: END RANGES
+
+// RUN: %sourcekitd-test -req=related-idents -pos=39:19 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK9 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=41:20 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK9 %s
+// RUN: %sourcekitd-test -req=related-idents -pos=43:11 %s -- -module-name related_idents %s | %FileCheck -check-prefix=CHECK9 %s
+// CHECK9:      START RANGES
+// CHECK9-NEXT: 39:19 - 1
+// CHECK9-NEXT: 41:20 - 1
+// CHECK9-NEXT: 43:11 - 1
+// CHECK9-NEXT: END RANGES

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1838,15 +1838,27 @@ class RelatedIdScanner : public SourceEntityWalker {
   explicit RelatedIdScanner(SourceFile &SrcFile, unsigned BufferID,
                             ValueDecl *D,
       llvm::SmallVectorImpl<std::pair<unsigned, unsigned>> &Ranges)
-    : Dcl(D), Ranges(Ranges),
-      SourceMgr(SrcFile.getASTContext().SourceMgr),
+    : Ranges(Ranges), 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:
+      //   case .first(let x), .second(let x)
+      //     fallthrough
+      //   case .third(let x)
+      //     print(x)
+      Dcl = V->getCanonicalVarDecl();
+    } else {
+      Dcl = D;
+    }
   }
 
 private:
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
     if (Cancelled)
       return false;
+    if (auto *V = dyn_cast<VarDecl>(D))
+        D = V->getCanonicalVarDecl();
     if (D == Dcl)
       return passId(Range);
     return true;
@@ -1856,8 +1868,12 @@ class RelatedIdScanner : public SourceEntityWalker {
                           ReferenceMetaData Data) override {
     if (Cancelled)
       return false;
-    if (CtorTyRef)
+
+    if (auto *V = dyn_cast<VarDecl>(D))
+      D = V->getCanonicalVarDecl();
+    else if (CtorTyRef)
       D = CtorTyRef;
+
     if (D == Dcl)
       return passId(Range);
     return true;
@@ -1933,7 +1949,13 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
           return;
 
         RelatedIdScanner Scanner(SrcFile, BufferID, VD, Ranges);
-        if (DeclContext *LocalDC = VD->getDeclContext()->getLocalContext()) {
+
+        if (auto *Case = getCaseStmtOfCanonicalVar(VD)) {
+          Scanner.walk(Case);
+          while ((Case = Case->getFallthroughDest().getPtrOrNull())) {
+            Scanner.walk(Case);
+          }
+        } else if (DeclContext *LocalDC = VD->getDeclContext()->getLocalContext()) {
           Scanner.walk(LocalDC);
         } else {
           Scanner.walk(SrcFile);
@@ -1956,6 +1978,16 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
       LOG_WARN_FUNC("related idents failed: " << Error);
       Receiver({});
     }
+
+    static CaseStmt *getCaseStmtOfCanonicalVar(Decl *D) {
+      assert(D);
+      if (auto *VD = dyn_cast<VarDecl>(D)) {
+        if (auto *Canonical = VD->getCanonicalVarDecl()) {
+          return dyn_cast_or_null<CaseStmt>(Canonical->getRecursiveParentPatternStmt());
+        }
+      }
+      return nullptr;
+    }
   };
 
   auto Consumer = std::make_shared<RelatedIdConsumer>(Offset, Receiver, Invok);