[Index] Handle shorthand if let/closure captures in local rename

Update rename to pull the outermost-declaration so that references are correctly found.

Rather than keeping suppressed locations in the current parent, keep
them for the whole index. Local rename starts the lookup from the
innermost context it can, which could be a closure. In that case there
is no parent decl on the stack and thus no where to store the locations
to suppress. We could have a specific store for this case, but there
shouldn't be that many of these and they're relatively cheap to store
anyway.

Resolves rdar://104568539.

Author: bnbarham

include/swift/IDE/Utils.h

@@ -161,7 +161,8 @@ struct ResolvedCursorInfo {
     /// that names both the newly declared variable and the referenced variable
     /// by the same identifier in the source text. This includes shorthand
     /// closure captures (`[foo]`) and shorthand if captures
-    /// (`if let foo {`).
+    /// (`if let foo {`). Ordered from innermost to outermost shadows.
+    ///
     /// Decls that are shadowed using shorthand syntax should be reported as
     /// additional cursor info results.
     SmallVector<ValueDecl *> ShorthandShadowedDecls;

lib/IDE/CursorInfo.cpp

@@ -112,7 +112,7 @@ class NodeFinder : ASTWalker {
   std::unique_ptr<NodeFinderResult> takeResult() { return std::move(Result); }
 
   /// Get the declarations that \p ShadowingDecl shadows using shorthand shadow
-  /// syntax.
+  /// syntax. Ordered from innermost to outermost shadows.
   SmallVector<ValueDecl *, 2>
   getShorthandShadowedDecls(ValueDecl *ShadowingDecl) {
     SmallVector<ValueDecl *, 2> Result;

lib/IDE/IDERequests.cpp

@@ -64,7 +64,8 @@ class CursorInfoResolver : public SourceEntityWalker {
   llvm::SmallVector<Expr*, 8> ExprStack;
   /// If a decl shadows another decl using shorthand syntax (`[foo]` or
   /// `if let foo {`), this maps the re-declared variable to the one that is
-  /// being shadowed.
+  /// being shadowed. Ordered from innermost to outermost shadows.
+  ///
   /// The transitive closure of shorthand shadowed decls should be reported as
   /// additional results in cursor info.
   llvm::DenseMap<ValueDecl *, ValueDecl *> ShorthandShadowedDecls;
@@ -180,17 +181,14 @@ ResolvedCursorInfo CursorInfoResolver::resolve(SourceLoc Loc) {
   walk(SrcFile);
 
   if (auto ValueRefInfo = dyn_cast<ResolvedValueRefCursorInfo>(&CursorInfo)) {
-    if (!ValueRefInfo->isRef()) {
-      SmallVector<ValueDecl *> ShadowedDecls;
-      // If we have a definition, add any decls that it potentially shadows
-      auto ShorthandShadowedDecl =
-          ShorthandShadowedDecls[ValueRefInfo->getValueD()];
-      while (ShorthandShadowedDecl) {
-        ShadowedDecls.push_back(ShorthandShadowedDecl);
-        ShorthandShadowedDecl = ShorthandShadowedDecls[ShorthandShadowedDecl];
-      }
-      ValueRefInfo->setShorthandShadowedDecls(ShadowedDecls);
+    SmallVector<ValueDecl *> ShadowedDecls;
+    auto ShorthandShadowedDecl =
+        ShorthandShadowedDecls[ValueRefInfo->getValueD()];
+    while (ShorthandShadowedDecl) {
+      ShadowedDecls.push_back(ShorthandShadowedDecl);
+      ShorthandShadowedDecl = ShorthandShadowedDecls[ShorthandShadowedDecl];
     }
+    ValueRefInfo->setShorthandShadowedDecls(ShadowedDecls);
   }
 
   return CursorInfo;

lib/Index/Index.cpp

@@ -432,7 +432,6 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     SymbolInfo SymInfo;
     SymbolRoleSet Roles;
     SmallVector<IndexedWitness, 6> ExplicitWitnesses;
-    SmallVector<SourceLoc, 6> RefsToSuppress;
   };
   SmallVector<Entity, 6> EntitiesStack;
   SmallVector<Expr *, 8> ExprStack;
@@ -449,6 +448,9 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
   StringScratchSpace stringStorage;
   ContainerTracker Containers;
 
+  // Already handled references that should be suppressed if found later.
+  llvm::DenseSet<SourceLoc> RefsToSuppress;
+
   // 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.
@@ -832,7 +834,7 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
                           ReferenceMetaData Data) override {
     SourceLoc Loc = Range.getStart();
 
-    if (isRepressed(Loc) || Loc.isInvalid())
+    if (Loc.isInvalid() || isSuppressed(Loc))
       return true;
 
     // Dig back to the original captured variable
@@ -912,17 +914,13 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     });
   }
 
-  void repressRefAtLoc(SourceLoc Loc) {
+  void suppressRefAtLoc(SourceLoc Loc) {
     if (Loc.isInvalid()) return;
-    assert(!EntitiesStack.empty());
-    EntitiesStack.back().RefsToSuppress.push_back(Loc);
+    RefsToSuppress.insert(Loc);
   }
 
-  bool isRepressed(SourceLoc Loc) const {
-    if (EntitiesStack.empty() || Loc.isInvalid())
-      return false;
-    auto &Suppressed = EntitiesStack.back().RefsToSuppress;
-    return std::find(Suppressed.begin(), Suppressed.end(), Loc) != Suppressed.end();
+  bool isSuppressed(SourceLoc Loc) const {
+    return Loc.isValid() && RefsToSuppress.contains(Loc);
   }
 
   Expr *getContainingExpr(size_t index) const {
@@ -1253,7 +1251,7 @@ bool IndexSwiftASTWalker::startEntity(Decl *D, IndexSymbol &Info, bool IsRef) {
         if (!handleWitnesses(D, explicitWitnesses))
           return false;
       }
-      EntitiesStack.push_back({D, Info.symInfo, Info.roles, std::move(explicitWitnesses), {}});
+      EntitiesStack.push_back({D, Info.symInfo, Info.roles, std::move(explicitWitnesses)});
       return true;
     }
   }
@@ -1325,7 +1323,7 @@ bool IndexSwiftASTWalker::reportRelatedRef(ValueDecl *D, SourceLoc Loc, bool isI
     Info.roles |= (unsigned)SymbolRole::Implicit;
 
   // don't report this ref again when visitDeclReference reports it
-  repressRefAtLoc(Loc);
+  suppressRefAtLoc(Loc);
 
   if (!reportRef(D, Loc, Info, None)) {
     Cancelled = true;
@@ -1491,7 +1489,7 @@ bool IndexSwiftASTWalker::report(ValueDecl *D) {
     // 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);
+    suppressRefAtLoc(loc);
 
     // Skip the definition of a shadowed decl
     return true;

lib/Refactoring/Refactoring.cpp

@@ -912,6 +912,11 @@ bool RefactoringActionLocalRename::performChange() {
   auto ValueRefCursorInfo = dyn_cast<ResolvedValueRefCursorInfo>(&CursorInfo);
   if (ValueRefCursorInfo && ValueRefCursorInfo->getValueD()) {
     ValueDecl *VD = ValueRefCursorInfo->typeOrValue();
+    // The index always uses the outermost shadow for references
+    if (!ValueRefCursorInfo->getShorthandShadowedDecls().empty()) {
+      VD = ValueRefCursorInfo->getShorthandShadowedDecls().back();
+    }
+
     SmallVector<DeclContext *, 8> Scopes;
 
     Optional<RenameRefInfo> RefInfo;

test/refactoring/rename/shorthand_shadow.swift

@@ -0,0 +1,16 @@
+// Local rename starts from the `DeclContext` of the renamed `Decl`. For
+// closures that means we have no parents, so check that case specifically.
+
+func renameInClosure() {
+  // RUN: %refactor -rename -dump-text -source-filename %s -pos=%(line+2):10 -new-name=renamed | %FileCheck %s
+  // CHECK: shorthand_shadow.swift [[# @LINE+1]]:10 -> [[# @LINE+1]]:18
+  _ = { (toRename: Int?) in
+    // RUN: %refactor -rename -dump-text -source-filename %s -pos=%(line+2):12 -new-name=renamed | %FileCheck %s
+    // CHECK: shorthand_shadow.swift [[# @LINE+1]]:12 -> [[# @LINE+1]]:20
+    if let toRename {
+      // RUN: %refactor -rename -dump-text -source-filename %s -pos=%(line+2):11 -new-name=renamed | %FileCheck %s
+      // CHECK: shorthand_shadow.swift [[# @LINE+1]]:11 -> [[# @LINE+1]]:19
+      _ = toRename
+    }
+  }
+}

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1172,6 +1172,7 @@ static bool passCursorInfoForDecl(
     });
     return false;
   }
+
   if (MainInfo.VD != OrigInfo.VD && !OrigInfo.Unavailable) {
     CursorSymbolInfo &CtorSymbol = Symbols.emplace_back();
     if (auto Err =
@@ -1182,16 +1183,21 @@ static bool passCursorInfoForDecl(
       Symbols.pop_back();
     }
   }
-  for (auto D : Info.getShorthandShadowedDecls()) {
-    CursorSymbolInfo &SymbolInfo = Symbols.emplace_back();
-    DeclInfo DInfo(D, Type(), /*IsRef=*/true, /*IsDynamic=*/false,
-                   ArrayRef<NominalTypeDecl *>(), Invoc);
-    if (auto Err =
-            fillSymbolInfo(SymbolInfo, DInfo, Info.getLoc(), AddSymbolGraph,
-                           Lang, Invoc, PreviousSnaps, Allocator)) {
-      // Ignore but make sure to remove the partially-filled symbol
-      llvm::handleAllErrors(std::move(Err), [](const llvm::StringError &E) {});
-      Symbols.pop_back();
+
+  // Add in shadowed declarations if on a decl. For references just go to the
+  // actual declaration.
+  if (!Info.isRef()) {
+    for (auto D : Info.getShorthandShadowedDecls()) {
+      CursorSymbolInfo &SymbolInfo = Symbols.emplace_back();
+      DeclInfo DInfo(D, Type(), /*IsRef=*/true, /*IsDynamic=*/false,
+                     ArrayRef<NominalTypeDecl *>(), Invoc);
+      if (auto Err =
+          fillSymbolInfo(SymbolInfo, DInfo, Info.getLoc(), AddSymbolGraph,
+                         Lang, Invoc, PreviousSnaps, Allocator)) {
+        // Ignore but make sure to remove the partially-filled symbol
+        llvm::handleAllErrors(std::move(Err), [](const llvm::StringError &E) {});
+        Symbols.pop_back();
+      }
     }
   }