[clangd] Rename: Allow multiple canonical declarations [1/3]

[tom-anders]

Summary

Note that locateDeclAt() now does *not* canonicalize the declaration anymore,
since that would now lead to "AmbiguousSymbol" when there's a single declaration under the cursor that has multiple canonical declarations.
Instead, only canonicalize *after* checking how many decls we've found under the cursor.

This will allow renaming virtual methods with size_overridden_methods() > 1 in a follow-up commit.

[tom-anders]

clang-format

clang-tools-extra/clangd/refactor/Rename.cpp

 //    template <typename T>
 //    bool Foo<T, 0> = true; (2)
 //
 //    template <>
 //    bool Foo<int, 0> = true; (3)
 //
 // Here, both partial (2) and full (3) specializations are canonicalized to (1)
 // which ensures all three of them are renamed.
-const NamedDecl *canonicalRenameDecl(const NamedDecl *D) {
+llvm::DenseSet<const NamedDecl *> canonicalRenameDecls(const NamedDecl *D) {
   if (const auto *VarTemplate = dyn_cast<VarTemplateSpecializationDecl>(D))
-    return canonicalRenameDecl(
+    return canonicalRenameDecls(
         VarTemplate->getSpecializedTemplate()->getTemplatedDecl());
   if (const auto *Template = dyn_cast<TemplateDecl>(D))
     if (const NamedDecl *TemplatedDecl = Template->getTemplatedDecl())
-      return canonicalRenameDecl(TemplatedDecl);
+      return canonicalRenameDecls(TemplatedDecl);
   if (const auto *ClassTemplateSpecialization =
           dyn_cast<ClassTemplateSpecializationDecl>(D))
-    return canonicalRenameDecl(
+    return canonicalRenameDecls(
         ClassTemplateSpecialization->getSpecializedTemplate()
             ->getTemplatedDecl());
   if (const auto *Method = dyn_cast<CXXMethodDecl>(D)) {
     if (Method->getDeclKind() == Decl::Kind::CXXConstructor ||
         Method->getDeclKind() == Decl::Kind::CXXDestructor)
-      return canonicalRenameDecl(Method->getParent());
+      return canonicalRenameDecls(Method->getParent());
     if (const FunctionDecl *InstantiatedMethod =
             Method->getInstantiatedFromMemberFunction())
-      return canonicalRenameDecl(InstantiatedMethod);
+      return canonicalRenameDecls(InstantiatedMethod);
     // FIXME(kirillbobyrev): For virtual methods with
     // size_overridden_methods() > 1, this will not rename all functions it
     // overrides, because this code assumes there is a single canonical
     // declaration.
     if (Method->isVirtual() && Method->size_overridden_methods())
-      return canonicalRenameDecl(*Method->overridden_methods().begin());
+      return {canonicalRenameDecls(*Method->overridden_methods().begin())};
   }
   if (const auto *Function = dyn_cast<FunctionDecl>(D))
     if (const FunctionTemplateDecl *Template = Function->getPrimaryTemplate())
-      return canonicalRenameDecl(Template);
+      return canonicalRenameDecls(Template);
   if (const auto *Field = dyn_cast<FieldDecl>(D)) {
     // This is a hacky way to do something like
     // CXXMethodDecl::getInstantiatedFromMemberFunction for the field because
     // Clang AST does not store relevant information about the field that is
     // instantiated.
     const auto *FieldParent =
         dyn_cast_or_null<CXXRecordDecl>(Field->getParent());
     if (!FieldParent)
-      return Field->getCanonicalDecl();
+      return {Field->getCanonicalDecl()};
     FieldParent = FieldParent->getTemplateInstantiationPattern();
     // Field is not instantiation.
     if (!FieldParent || Field->getParent() == FieldParent)
-      return Field->getCanonicalDecl();
+      return {Field->getCanonicalDecl()};
     for (const FieldDecl *Candidate : FieldParent->fields())
       if (Field->getDeclName() == Candidate->getDeclName())
-        return Candidate->getCanonicalDecl();
+        return {Candidate->getCanonicalDecl()};
     elog("FieldParent should have field with the same name as Field.");
   }
   if (const auto *VD = dyn_cast<VarDecl>(D)) {
     if (const VarDecl *OriginalVD = VD->getInstantiatedFromStaticDataMember())
-      return canonicalRenameDecl(OriginalVD);
+      return canonicalRenameDecls(OriginalVD);
   }
   if (const auto *UD = dyn_cast<UsingShadowDecl>(D)) {
     if (const auto *TargetDecl = UD->getTargetDecl())
-      return canonicalRenameDecl(TargetDecl);
+      return canonicalRenameDecls(TargetDecl);
   }
-  return dyn_cast<NamedDecl>(D->getCanonicalDecl());
+  return {dyn_cast<NamedDecl>(D->getCanonicalDecl())};
 }
 
 llvm::DenseSet<const NamedDecl *> locateDeclAt(ParsedAST &AST,
                                                SourceLocation TokenStartLoc) {
   unsigned Offset =
       AST.getSourceManager().getDecomposedSpellingLoc(TokenStartLoc).second;
 
   SelectionTree Selection = SelectionTree::createRight(
       AST.getASTContext(), AST.getTokens(), Offset, Offset);
   const SelectionTree::Node *SelectedNode = Selection.commonAncestor();
   if (!SelectedNode)
     return {};
 
   llvm::DenseSet<const NamedDecl *> Result;
   for (const NamedDecl *D :
        targetDecl(SelectedNode->ASTNode,
                   DeclRelation::Alias | DeclRelation::TemplatePattern,
                   AST.getHeuristicResolver())) {
-    Result.insert(canonicalRenameDecl(D));
+    Result.insert(D);
   }
   return Result;
 }
 
 void filterRenameTargets(llvm::DenseSet<const NamedDecl *> &Decls) {
   // For something like
   //     namespace ns { void foo(); }
   //     void bar() { using ns::f^oo; foo(); }
llvm::Error makeError(ReasonToReject Reason) {
   };
   return error("Cannot rename symbol: {0}", Message(Reason));
 }
 
 // Return all rename occurrences in the main file.
 std::vector<SourceLocation> findOccurrencesWithinFile(ParsedAST &AST,
                                                       const NamedDecl &ND) {
   trace::Span Tracer("FindOccurrencesWithinFile");
-  assert(canonicalRenameDecl(&ND) == &ND &&
+  assert(canonicalRenameDecls(&ND).contains(&ND) &&
          "ND should be already canonicalized.");
 
   std::vector<SourceLocation> Results;
   for (Decl *TopLevelDecl : AST.getLocalTopLevelDecls()) {
     findExplicitReferences(
         TopLevelDecl,
         [&](ReferenceLoc Ref) {
           if (Ref.Targets.empty())
             return;
           for (const auto *Target : Ref.Targets) {
-            if (canonicalRenameDecl(Target) == &ND) {
+            if (canonicalRenameDecls(Target).contains(&ND)) {
               Results.push_back(Ref.NameLoc);
               return;
             }
           }
         },
         AST.getHeuristicResolver());
   }
 
llvm::Optional<InvalidName> checkName(const NamedDecl &RenameDecl,
   }
   if (Result)
     InvalidNameMetric.record(1, toString(Result->K));
   return Result;
 }
 
 // AST-based rename, it renames all occurrences in the main file.
 llvm::Expected<tooling::Replacements>
-renameWithinFile(ParsedAST &AST, const NamedDecl &RenameDecl,
+renameWithinFile(ParsedAST &AST,
+                 const llvm::DenseSet<const NamedDecl *> &RenameDecls,
                  llvm::StringRef NewName) {
   trace::Span Tracer("RenameWithinFile");
   const SourceManager &SM = AST.getSourceManager();
 
   tooling::Replacements FilteredChanges;
-  for (SourceLocation Loc : findOccurrencesWithinFile(AST, RenameDecl)) {
+  for (const auto *RenameDecl : RenameDecls)
+    for (SourceLocation Loc : findOccurrencesWithinFile(AST, *RenameDecl)) {
       SourceLocation RenameLoc = Loc;
       // We don't rename in any macro bodies, but we allow rename the symbol
       // spelled in a top-level macro argument in the main file.
       if (RenameLoc.isMacroID()) {
         if (isInMacroBody(SM, RenameLoc))
           continue;
         RenameLoc = SM.getSpellingLoc(Loc);
       }
       // Filter out locations not from main file.
       // We traverse only main file decls, but locations could come from an
       // non-preamble #include file e.g.
       //   void test() {
       //     int f^oo;
       //     #include "use_foo.inc"
       //   }
       if (!isInsideMainFile(RenameLoc, SM))
         continue;
       if (auto Err = FilteredChanges.add(tooling::Replacement(
               SM, CharSourceRange::getTokenRange(RenameLoc), NewName)))
         return std::move(Err);
     }
   return FilteredChanges;
 }
 
 Range toRange(const SymbolLocation &L) {
   Range R;
   R.start.line = L.Start.line();
   R.start.character = L.Start.column();
   R.end.line = L.End.line();
Index.relations(Req, [&](const SymbolID &, const Symbol &Override) {
         Pending.insert(Override.ID);
     });
   }
 }
 
 // Return all rename occurrences (using the index) outside of the main file,
 // grouped by the absolute file path.
 llvm::Expected<llvm::StringMap<std::vector<Range>>>
-findOccurrencesOutsideFile(const NamedDecl &RenameDecl,
+findOccurrencesOutsideFile(const llvm::DenseSet<const NamedDecl *> &RenameDecls,
                            llvm::StringRef MainFile, const SymbolIndex &Index,
                            size_t MaxLimitFiles) {
+  assert(!RenameDecls.empty() &&
+         "Should have at least one decl to rename here");
   trace::Span Tracer("FindOccurrencesOutsideFile");
   RefsRequest RQuest;
-  RQuest.IDs.insert(getSymbolID(&RenameDecl));
 
-  if (const auto *MethodDecl = llvm::dyn_cast<CXXMethodDecl>(&RenameDecl))
+  for (const NamedDecl *D : RenameDecls) {
+    RQuest.IDs.insert(getSymbolID(D));
+    if (const auto *MethodDecl = llvm::dyn_cast<CXXMethodDecl>(D))
       if (MethodDecl->isVirtual())
         insertTransitiveOverrides(*RQuest.IDs.begin(), RQuest.IDs, Index);
+  }
 
   // Absolute file path => rename occurrences in that file.
   llvm::StringMap<std::vector<Range>> AffectedFiles;
   bool HasMore = Index.refs(RQuest, [&](const Ref &R) {
     if (AffectedFiles.size() >= MaxLimitFiles)
       return;
     if ((R.Kind & RefKind::Spelled) == RefKind::Unknown)
       return;
     if (auto RefFilePath = filePath(R.Location, /*HintFilePath=*/MainFile)) {
       if (!pathEqual(*RefFilePath, MainFile))
         AffectedFiles[*RefFilePath].push_back(toRange(R.Location));
     }
   });
 
   if (AffectedFiles.size() >= MaxLimitFiles)
     return error("The number of affected files exceeds the max limit {0}",
                  MaxLimitFiles);
   if (HasMore)
     return error("The symbol {0} has too many occurrences",
-                 RenameDecl.getQualifiedNameAsString());
+                 (*RenameDecls.begin())->getQualifiedNameAsString());
   // Sort and deduplicate the results, in case that index returns duplications.
   for (auto &FileAndOccurrences : AffectedFiles) {
     auto &Ranges = FileAndOccurrences.getValue();
     llvm::sort(Ranges);
     Ranges.erase(std::unique(Ranges.begin(), Ranges.end()), Ranges.end());
 
     SPAN_ATTACH(Tracer, FileAndOccurrences.first(),
                 static_cast<int64_t>(Ranges.size()));

 //
 // Clangd builds a dynamic index for all opened files on top of the static
 // index of the whole codebase. Dynamic index is up-to-date (respects dirty
 // buffers) as long as clangd finishes processing opened files, while static
 // index (background index) is relatively stale. We choose the dirty buffers
 // as the file content we rename on, and fallback to file content on disk if
 // there is no dirty buffer.
 llvm::Expected<FileEdits>
-renameOutsideFile(const NamedDecl &RenameDecl, llvm::StringRef MainFilePath,
-                  llvm::StringRef NewName, const SymbolIndex &Index,
-                  size_t MaxLimitFiles, llvm::vfs::FileSystem &FS) {
+renameOutsideFile(const llvm::DenseSet<const NamedDecl *> &RenameDecls,
+                  llvm::StringRef MainFilePath, llvm::StringRef NewName,
+                  const SymbolIndex &Index, size_t MaxLimitFiles,
+                  llvm::vfs::FileSystem &FS) {
   trace::Span Tracer("RenameOutsideFile");
-  auto AffectedFiles = findOccurrencesOutsideFile(RenameDecl, MainFilePath,
+  auto AffectedFiles = findOccurrencesOutsideFile(RenameDecls, MainFilePath,
                                                   Index, MaxLimitFiles);
   if (!AffectedFiles)
     return AffectedFiles.takeError();
   FileEdits Results;
+  for (const auto *RenameDecl : RenameDecls)
     for (auto &FileAndOccurrences : *AffectedFiles) {
       llvm::StringRef FilePath = FileAndOccurrences.first();
 
       auto ExpBuffer = FS.getBufferForFile(FilePath);
       if (!ExpBuffer) {
         elog("Fail to read file content: Fail to open file {0}: {1}", FilePath,
              ExpBuffer.getError().message());
         continue;
       }
 
       auto AffectedFileCode = (*ExpBuffer)->getBuffer();
-    auto RenameRanges =
-        adjustRenameRanges(AffectedFileCode, RenameDecl.getNameAsString(),
-                           std::move(FileAndOccurrences.second),
-                           RenameDecl.getASTContext().getLangOpts());
+      auto RenameRanges = adjustRenameRanges(
+          AffectedFileCode, RenameDecl->getNameAsString(),
+          FileAndOccurrences.second, RenameDecl->getASTContext().getLangOpts());
       if (!RenameRanges) {
         // Our heuristics fails to adjust rename ranges to the current state of
         // the file, it is most likely the index is stale, so we give up the
         // entire rename.
         return error("Index results don't match the content of file {0} "
                      "(the index may be stale)",
                      FilePath);
       }
       auto RenameEdit =
           buildRenameEdit(FilePath, AffectedFileCode, *RenameRanges, NewName);
       if (!RenameEdit)
         return error("failed to rename in file {0}: {1}", FilePath,
                      RenameEdit.takeError());
       if (!RenameEdit->Replacements.empty())
         Results.insert({FilePath, std::move(*RenameEdit)});
     }
   return Results;
 }
 
 // A simple edit is either changing line or column, but not both.
 bool impliesSimpleEdit(const Position &LHS, const Position &RHS) {
   return LHS.line == RHS.line || LHS.character == RHS.character;
 }
 
if (locateMacroAt(*IdentifierToken, AST.getPreprocessor()))
     return makeError(ReasonToReject::UnsupportedSymbol);
 
   auto DeclsUnderCursor = locateDeclAt(AST, IdentifierToken->location());
   filterRenameTargets(DeclsUnderCursor);
   if (DeclsUnderCursor.empty())
     return makeError(ReasonToReject::NoSymbolFound);
   if (DeclsUnderCursor.size() > 1)
     return makeError(ReasonToReject::AmbiguousSymbol);
-  const auto &RenameDecl = **DeclsUnderCursor.begin();
-  const auto *ID = RenameDecl.getIdentifier();
+
+  llvm::DenseSet<const NamedDecl *> RenameDecls;
+  for (const auto *D : DeclsUnderCursor) {
+    auto CD = canonicalRenameDecls(D);
+    RenameDecls.insert(CD.begin(), CD.end());
+  }
+
+  for (const NamedDecl *D : RenameDecls) {
+    auto *ID = D->getIdentifier();
     if (!ID)
       return makeError(ReasonToReject::UnsupportedSymbol);
     if (ID->getName() == RInputs.NewName)
       return makeError(ReasonToReject::SameName);
-  auto Invalid = checkName(RenameDecl, RInputs.NewName);
+
+    auto Invalid = checkName(*D, RInputs.NewName);
     if (Invalid)
       return makeError(std::move(*Invalid));
 
-  auto Reject =
-      renameable(RenameDecl, RInputs.MainFilePath, RInputs.Index, Opts);
+    auto Reject = renameable(*D, RInputs.MainFilePath, RInputs.Index, Opts);
     if (Reject)
       return makeError(*Reject);
+  }
 
   // We have two implementations of the rename:
   //   - AST-based rename: used for renaming local symbols, e.g. variables
   //     defined in a function body;
   //   - index-based rename: used for renaming non-local symbols, and not
   //     feasible for local symbols (as by design our index don't index these
   //     symbols by design;
   // To make cross-file rename work for local symbol, we use a hybrid solution:
   //   - run AST-based rename on the main file;
   //   - run index-based rename on other affected files;
-  auto MainFileRenameEdit = renameWithinFile(AST, RenameDecl, RInputs.NewName);
+  auto MainFileRenameEdit = renameWithinFile(AST, RenameDecls, RInputs.NewName);
   if (!MainFileRenameEdit)
     return MainFileRenameEdit.takeError();
 
   // Check the rename-triggering location is actually being renamed.
   // This is a robustness check to avoid surprising rename results -- if the
   // the triggering location is not actually the name of the node we identified
   // (e.g. for broken code), then rename is likely not what users expect, so we
   // reject this kind of rename.
if (llvm::none_of(
     return makeError(ReasonToReject::NoSymbolFound);
   }
   RenameResult Result;
   Result.Target = CurrentIdentifier;
   Edit MainFileEdits = Edit(MainFileCode, std::move(*MainFileRenameEdit));
   for (const TextEdit &TE : MainFileEdits.asTextEdits())
     Result.LocalChanges.push_back(TE.range);
 
-  // return the main file edit if this is a within-file rename or the symbol
-  // being renamed is function local.
-  if (RenameDecl.getParentFunctionOrMethod()) {
+  // return the main file edit if this is a within-file rename or the symbols
+  // being renamed are all function local.
+  auto IsWithinFileRename = [](const NamedDecl *D) {
+    return D->getParentFunctionOrMethod() != nullptr;
+  };
+  if (llvm::all_of(RenameDecls, IsWithinFileRename)) {
     Result.GlobalChanges = FileEdits(
         {std::make_pair(RInputs.MainFilePath, std::move(MainFileEdits))});
     return Result;
   }
 
   // If the index is nullptr, we don't know the completeness of the result, so
   // we don't populate the field GlobalChanges.
   if (!RInputs.Index) {
     assert(Result.GlobalChanges.empty());
     return Result;
   }
 
   auto OtherFilesEdits = renameOutsideFile(
-      RenameDecl, RInputs.MainFilePath, RInputs.NewName, *RInputs.Index,
+      RenameDecls, RInputs.MainFilePath, RInputs.NewName, *RInputs.Index,
       Opts.LimitFiles == 0 ? std::numeric_limits<size_t>::max()
                            : Opts.LimitFiles,
       *RInputs.FS);
   if (!OtherFilesEdits)
     return OtherFilesEdits.takeError();
   Result.GlobalChanges = *OtherFilesEdits;
   // Attach the rename edits for the main file.
   Result.GlobalChanges.try_emplace(RInputs.MainFilePath,