[5.8][Index] Resolve any shadowed references to the decl they shadow

Skip the shadowing decl and instead reference the top-most shadowed decl
(ie. follow any shadowed decls up the chain).

Also fix `getShorthandShadows` missing references in `MemberDeclRefExpr`
by using `getReferencedDecl` instead of assuming it is a `DeclRefExpr`.

Resolves rdar://99730146, rdar://100006415, and rdar://103449038.

(cherry picked from commit 7b4ab8d)

Author: bnbarham

lib/IDE/IDETypeChecking.cpp

@@ -918,35 +918,30 @@ SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
 swift::getShorthandShadows(CaptureListExpr *CaptureList, DeclContext *DC) {
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1> Result;
   for (auto Capture : CaptureList->getCaptureList()) {
-    if (Capture.PBD->getPatternList().size() != 1) {
+    if (Capture.PBD->getPatternList().size() != 1)
       continue;
-    }
+
     Expr *Init = Capture.PBD->getInit(0);
-    if (!Init) {
-      continue;
-    }
-    if (auto UDRE = dyn_cast<UnresolvedDeclRefExpr>(Init)) {
-      if (DC) {
-        Init = resolveDeclRefExpr(UDRE, DC, /*replaceInvalidRefsWithErrors=*/false);
-      }
-    }
-    auto *DRE = dyn_cast_or_null<DeclRefExpr>(Init);
-    if (!DRE) {
+    if (!Init)
       continue;
-    }
 
     auto DeclaredVar = Capture.getVar();
-    if (DeclaredVar->getLoc() != DRE->getLoc()) {
+    if (DeclaredVar->getLoc() != Init->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;
+    if (auto UDRE = dyn_cast<UnresolvedDeclRefExpr>(Init)) {
+      if (DC) {
+        Init = resolveDeclRefExpr(UDRE, DC, /*replaceInvalidRefsWithErrors=*/false);
+      }
     }
 
+    auto *ReferencedVar = Init->getReferencedDecl().getDecl();
+    if (!ReferencedVar)
+      continue;
+
     assert(DeclaredVar->getName() == ReferencedVar->getName());
 
     Result.emplace_back(std::make_pair(DeclaredVar, ReferencedVar));
@@ -958,28 +953,23 @@ SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
 swift::getShorthandShadows(LabeledConditionalStmt *CondStmt, DeclContext *DC) {
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1> Result;
   for (const StmtConditionElement &Cond : CondStmt->getCond()) {
-    if (Cond.getKind() != StmtConditionElement::CK_PatternBinding) {
+    if (Cond.getKind() != StmtConditionElement::CK_PatternBinding)
       continue;
-    }
+
     Expr *Init = Cond.getInitializer();
     if (auto UDRE = dyn_cast<UnresolvedDeclRefExpr>(Init)) {
       if (DC) {
         Init = resolveDeclRefExpr(UDRE, DC, /*replaceInvalidRefsWithErrors=*/false);
       }
     }
-    auto InitDeclRef = dyn_cast<DeclRefExpr>(Init);
-    if (!InitDeclRef) {
-      continue;
-    }
-    auto ReferencedVar = dyn_cast_or_null<VarDecl>(InitDeclRef->getDecl());
-    if (!ReferencedVar) {
+
+    auto ReferencedVar = Init->getReferencedDecl().getDecl();
+    if (!ReferencedVar)
       continue;
-    }
 
     Cond.getPattern()->forEachVariable([&](VarDecl *DeclaredVar) {
-      if (DeclaredVar->getLoc() != Init->getLoc()) {
+      if (DeclaredVar->getLoc() != Init->getLoc())
         return;
-      }
       assert(DeclaredVar->getName() == ReferencedVar->getName());
       Result.emplace_back(std::make_pair(DeclaredVar, ReferencedVar));
     });

lib/Index/Index.cpp

@@ -22,6 +22,7 @@
 #include "swift/AST/Pattern.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/Stmt.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/AST/Types.h"
 #include "swift/AST/USRGeneration.h"
@@ -448,6 +449,11 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
   StringScratchSpace stringStorage;
   ContainerTracker Containers;
 
+  // Contains a mapping for captures of the form [x], from the declared "x"
+  // to the captured "x" in the enclosing scope. Also includes shorthand if
+  // let bindings.
+  llvm::DenseMap<ValueDecl *, ValueDecl *> sameNamedCaptures;
+
   bool getNameAndUSR(ValueDecl *D, ExtensionDecl *ExtD,
                      StringRef &name, StringRef &USR) {
     auto &result = nameAndUSRCache[ExtD ? (Decl*)ExtD : D];
@@ -555,6 +561,16 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     return false;
   }
 
+  ValueDecl *firstDecl(ValueDecl *D) {
+    while (true) {
+      auto captured = sameNamedCaptures.find(D);
+      if (captured == sameNamedCaptures.end())
+        break;
+      D = captured->second;
+    }
+    return D;
+  }
+
 public:
   IndexSwiftASTWalker(IndexDataConsumer &IdxConsumer, ASTContext &Ctx,
                       SourceFile *SF = nullptr)
@@ -701,6 +717,15 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
   bool walkToExprPre(Expr *E) override {
     if (Cancelled)
       return false;
+
+    // Record any same named captures/shorthand if let bindings so we can
+    // treat their references as references to the original decl.
+    if (auto *captureList = dyn_cast<CaptureListExpr>(E)) {
+      for (auto shadows : getShorthandShadows(captureList)) {
+        sameNamedCaptures[shadows.first] = shadows.second;
+      }
+    }
+
     ExprStack.push_back(E);
     Containers.activateContainersFor(E);
     handleMemberwiseInitRefs(E);
@@ -715,6 +740,20 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     return true;
   }
 
+  bool walkToStmtPre(Stmt *stmt) override {
+    if (Cancelled)
+      return false;
+
+    // Record any same named captures/shorthand if let bindings so we can
+    // treat their references as references to the original decl.
+    if (auto *condition = dyn_cast<LabeledConditionalStmt>(stmt)) {
+      for (auto shadows : getShorthandShadows(condition)) {
+        sameNamedCaptures[shadows.first] = shadows.second;
+      }
+    }
+    return true;
+  }
+
   bool walkToTypeReprPre(TypeRepr *T) override {
     if (Cancelled)
       return false;
@@ -796,6 +835,11 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     if (isRepressed(Loc) || Loc.isInvalid())
       return true;
 
+    // Dig back to the original captured variable
+    if (auto *VD = dyn_cast<VarDecl>(D)) {
+      D = firstDecl(D);
+    }
+
     IndexSymbol Info;
 
     if (Data.isImplicit)
@@ -1438,6 +1482,24 @@ bool IndexSwiftASTWalker::reportExtension(ExtensionDecl *D) {
 }
 
 bool IndexSwiftASTWalker::report(ValueDecl *D) {
+  auto *shadowedDecl = firstDecl(D);
+  if (D != shadowedDecl) {
+    // Report a reference to the shadowed decl
+    SourceLoc loc = D->getNameLoc();
+
+    IndexSymbol info;
+    if (!reportRef(shadowedDecl, loc, info, AccessKind::Read))
+      return false;
+
+    // Suppress the reference if there is any (it is implicit and hence
+    // already skipped in the shorthand if let case, but explicit in the
+    // captured case).
+    repressRefAtLoc(loc);
+
+    // Skip the definition of a shadowed decl
+    return true;
+  }
+
   if (startEntityDecl(D)) {
     // Pass accessors.
     if (auto StoreD = dyn_cast<AbstractStorageDecl>(D)) {

test/Index/index_shadow.swift

@@ -0,0 +1,60 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-ide-test -print-indexed-symbols -source-filename %s --include-locals | %FileCheck %s
+
+// The index will output references to the shadowed-declaration rather than
+// the one defined by the shorthand if-let or capture. It also skips
+// outputting the shadowing-definiiton since it would then have no references.
+
+struct ShadowedTest {
+  // CHECK: [[@LINE+1]]:7 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Def
+  let shadowedVar: Int?? = 1
+
+  func localShadowTest() {
+    // CHECK: [[@LINE+1]]:9 {{.*}} s:14swift_ide_test12ShadowedTestV011localShadowE0yyF0fD3VarL_SiSgvp {{.*}}Def
+    let localShadowedVar: Int? = 2
+
+    // CHECK-NOT: [[@LINE+2]]:12 {{.*}} localShadowedVar {{.*}}Def
+    // CHECK: [[@LINE+1]]:12 {{.*}} s:14swift_ide_test12ShadowedTestV011localShadowE0yyF0fD3VarL_SiSgvp {{.*}}Ref
+    if let localShadowedVar {
+      // CHECK: [[@LINE+1]]:11 {{.*}} s:14swift_ide_test12ShadowedTestV011localShadowE0yyF0fD3VarL_SiSgvp {{.*}}Ref
+      _ = localShadowedVar
+    }
+
+    // CHECK-NOT: [[@LINE+2]]:12 {{.*}} localShadowedVar {{.*}}Def
+    // CHECK: [[@LINE+1]]:12 {{.*}} s:14swift_ide_test12ShadowedTestV011localShadowE0yyF0fD3VarL_SiSgvp {{.*}}Ref
+    _ = { [localShadowedVar] in
+      // CHECK: [[@LINE+1]]:11 {{.*}} s:14swift_ide_test12ShadowedTestV011localShadowE0yyF0fD3VarL_SiSgvp {{.*}}Ref
+      _ = localShadowedVar
+    }
+  }
+
+  func shadowTest() {
+    // CHECK-NOT: [[@LINE+2]]:12 {{.*}} shadowedVar {{.*}}Def
+    // CHECK: [[@LINE+1]]:12 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+    if let shadowedVar {
+      // CHECK: [[@LINE+1]]:11 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+      _ = shadowedVar
+
+      // CHECK-NOT: [[@LINE+2]]:14 {{.*}} shadowedVar {{.*}}Def
+      // CHECK: [[@LINE+1]]:14 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+      if let shadowedVar {
+        // CHECK: [[@LINE+1]]:13 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+        _ = shadowedVar
+      }
+    }
+
+    // CHECK-NOT: [[@LINE+2]]:12 {{.*}} shadowedVar {{.*}}Def
+    // CHECK: [[@LINE+1]]:12 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+    _ = { [shadowedVar] in
+      // CHECK: [[@LINE+1]]:11 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+      _ = shadowedVar
+
+      // CHECK-NOT: [[@LINE+2]]:14 {{.*}} shadowedVar {{.*}}Def
+      // CHECK: [[@LINE+1]]:14 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+      _ = { [shadowedVar] in
+        // CHECK: [[@LINE+1]]:13 {{.*}} s:14swift_ide_test12ShadowedTestV11shadowedVarSiSgSgvp {{.*}}Ref
+        _ = shadowedVar
+      }
+    }
+  }
+}